MQE: shift responsibilities of Materializer (#12377)

#### What this PR does

This PR changes the responsibilities of MQE's `Materializer` type.

Previously, `Materializer.ConvertNodeToOperator` would materialize all
of the given node's children and pass it to `Node.OperatorFactory`, and
`OperatorFactory` would have no access to the `Materializer`. With the
existing set of query plan nodes, this was not an issue, but it is
problematic for some scenarios, for example:

* for remote execution, there's no point materializing operators for the
parts of the query plan that will be executed remotely
* for caching and splitting, there's also no point materializing
operators for cached results, but we'll then need to materialize the
child operator if there is a cache miss, or only materialize the child
operator for a subset of the time range

This PR also moves `Materializer` to the `planning` package to prevent
an import cycle (and seems like a more natural fit for there anyway).

Depending on which PR merges first, either this PR or
#12302 will need changes to account
for the changes in the other.

This PR has no user-visible impact and so I've chosen not to add a
changelog entry.

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

(none)

#### Checklist

- [n/a] Tests updated.
- [n/a] Documentation added.
- [n/a] `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.

Author: charleskorn

pkg/streamingpromql/engine.go

@@ -168,7 +168,7 @@ func (e *Engine) newQueryFromPlanner(ctx context.Context, queryable storage.Quer
 		EagerLoadSelectors:       e.eagerLoadSelectors,
 	}
 
-	materializer := NewMaterializer(operatorParams)
+	materializer := planning.NewMaterializer(operatorParams)
 	root, err := materializer.ConvertNodeToOperator(plan.Root, plan.TimeRange)
 	if err != nil {
 		return nil, err

pkg/streamingpromql/materialize.go

@@ -67,12 +67,13 @@ func (d *Duplicate) ResultType() (parser.ValueType, error) {
 	return d.Inner.ResultType()
 }
 
-func (d *Duplicate) OperatorFactory(children []types.Operator, _ types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
-	if len(children) != 1 {
-		return nil, fmt.Errorf("expected exactly 1 child for Duplicate, got %v", len(children))
+func (d *Duplicate) OperatorFactory(materializer *planning.Materializer, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
+	inner, err := materializer.ConvertNodeToOperator(d.Inner, timeRange)
+	if err != nil {
+		return nil, err
 	}
 
-	switch inner := children[0].(type) {
+	switch inner := inner.(type) {
 	case types.InstantVectorOperator:
 		return &InstantVectorDuplicationConsumerOperatorFactory{
 			Buffer: NewInstantVectorDuplicationBuffer(inner, params.MemoryConsumptionTracker),
@@ -82,7 +83,7 @@ func (d *Duplicate) OperatorFactory(children []types.Operator, _ types.QueryTime
 			Buffer: NewRangeVectorDuplicationBuffer(inner, params.MemoryConsumptionTracker),
 		}, nil
 	default:
-		return nil, fmt.Errorf("expected InstantVectorOperator as child of Duplicate, got %T", children[0])
+		return nil, fmt.Errorf("expected InstantVectorOperator or RangeVectorOperator as child of Duplicate, got %T", inner)
 	}
 }
 

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

@@ -100,64 +100,43 @@ func (a *AggregateExpression) ChildrenLabels() []string {
 	return []string{"expression", "parameter"}
 }
 
-func (a *AggregateExpression) OperatorFactory(children []types.Operator, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
-	if len(children) < 1 {
-		return nil, fmt.Errorf("expected at least 1 child for AggregateExpression, got %v", len(children))
-	}
-
-	inner, ok := children[0].(types.InstantVectorOperator)
-	if !ok {
-		return nil, fmt.Errorf("expected InstantVectorOperator as expression child of AggregateExpression, got %T", children[0])
+func (a *AggregateExpression) OperatorFactory(materializer *planning.Materializer, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
+	inner, err := materializer.ConvertNodeToInstantVectorOperator(a.Inner, timeRange)
+	if err != nil {
+		return nil, fmt.Errorf("could not create inner operator for AggregateExpression: %w", err)
 	}
 
 	var o types.InstantVectorOperator
 
 	switch a.Op {
 	case AGGREGATION_TOPK, AGGREGATION_BOTTOMK:
-		if len(children) != 2 {
-			return nil, fmt.Errorf("expected exactly 2 children for AggregateExpression with operation %s, got %v", a.Op.String(), len(children))
-		}
-
-		param, ok := children[1].(types.ScalarOperator)
-		if !ok {
-			return nil, fmt.Errorf("expected ScalarOperator as parameter child of AggregateExpression with operation %s, got %T", a.Op.String(), children[0])
+		param, err := materializer.ConvertNodeToScalarOperator(a.Param, timeRange)
+		if err != nil {
+			return nil, fmt.Errorf("could not create parameter operator for AggregateExpression %s: %w", a.Op.String(), err)
 		}
 
 		o = topkbottomk.New(inner, param, timeRange, a.Grouping, a.Without, a.Op == AGGREGATION_TOPK, params.MemoryConsumptionTracker, params.Annotations, a.ExpressionPosition.ToPrometheusType())
 
 	case AGGREGATION_QUANTILE:
-		if len(children) != 2 {
-			return nil, fmt.Errorf("expected exactly 2 children for AggregateExpression with operation %s, got %v", a.Op.String(), len(children))
-		}
-
-		param, ok := children[1].(types.ScalarOperator)
-		if !ok {
-			return nil, fmt.Errorf("expected ScalarOperator as parameter child of AggregateExpression with operation %s, got %T", a.Op.String(), children[0])
+		param, err := materializer.ConvertNodeToScalarOperator(a.Param, timeRange)
+		if err != nil {
+			return nil, fmt.Errorf("could not create parameter operator for AggregateExpression %s: %w", a.Op.String(), err)
 		}
 
-		var err error
 		o, err = aggregations.NewQuantileAggregation(inner, param, timeRange, a.Grouping, a.Without, params.MemoryConsumptionTracker, params.Annotations, a.ExpressionPosition.ToPrometheusType())
 		if err != nil {
 			return nil, err
 		}
 
 	case AGGREGATION_COUNT_VALUES:
-		if len(children) != 2 {
-			return nil, fmt.Errorf("expected exactly 2 children for AggregateExpression with operation %s, got %v", a.Op.String(), len(children))
-		}
-
-		param, ok := children[1].(types.StringOperator)
-		if !ok {
-			return nil, fmt.Errorf("expected StringOperator as parameter child of AggregateExpression with operation %s, got %T", a.Op.String(), children[0])
+		param, err := materializer.ConvertNodeToStringOperator(a.Param, timeRange)
+		if err != nil {
+			return nil, fmt.Errorf("could not create parameter operator for AggregateExpression %s: %w", a.Op.String(), err)
 		}
 
 		o = aggregations.NewCountValues(inner, param, timeRange, a.Grouping, a.Without, params.MemoryConsumptionTracker, a.ExpressionPosition.ToPrometheusType())
 
 	default:
-		if len(children) != 1 {
-			return nil, fmt.Errorf("expected exactly 1 child for AggregateExpression with operation %s, got %v", a.Op.String(), len(children))
-		}
-
 		itemType, ok := a.Op.ToItemType()
 		if !ok {
 			return nil, fmt.Errorf("unknown aggregation operation %s", a.Op.String())

pkg/streamingpromql/planning/core/aggregate_expression.go

@@ -117,25 +117,20 @@ func (b *BinaryExpression) ChildrenLabels() []string {
 	return []string{"LHS", "RHS"}
 }
 
-func (b *BinaryExpression) OperatorFactory(children []types.Operator, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
+func (b *BinaryExpression) OperatorFactory(materializer *planning.Materializer, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
 	op, ok := b.Op.ToItemType()
 	if !ok {
 		return nil, compat.NewNotSupportedError(fmt.Sprintf("'%v' binary expression", b.Op.String()))
 	}
 
-	if len(children) != 2 {
-		return nil, fmt.Errorf("expected exactly 2 children for BinaryExpression, got %v", len(children))
-	}
-
-	lhsVector, lhsScalar := b.getChildOperator(children[0])
-	rhsVector, rhsScalar := b.getChildOperator(children[1])
-
-	if lhsVector == nil && lhsScalar == nil {
-		return nil, fmt.Errorf("expected either InstantVectorOperator or ScalarOperator on left-hand side of BinaryExpression, got %T", children[0])
+	lhsVector, lhsScalar, err := b.getChildOperator(b.LHS, timeRange, materializer, "left")
+	if err != nil {
+		return nil, err
 	}
 
-	if rhsVector == nil && rhsScalar == nil {
-		return nil, fmt.Errorf("expected either InstantVectorOperator or ScalarOperator on right-hand side of BinaryExpression, got %T", children[1])
+	rhsVector, rhsScalar, err := b.getChildOperator(b.RHS, timeRange, materializer, "right")
+	if err != nil {
+		return nil, err
 	}
 
 	if lhsScalar != nil && rhsScalar != nil {
@@ -177,14 +172,19 @@ func (b *BinaryExpression) OperatorFactory(children []types.Operator, timeRange
 	return planning.NewSingleUseOperatorFactory(operators.NewDeduplicateAndMerge(o, params.MemoryConsumptionTracker)), nil
 }
 
-func (b *BinaryExpression) getChildOperator(child types.Operator) (types.InstantVectorOperator, types.ScalarOperator) {
-	switch child := child.(type) {
+func (b *BinaryExpression) getChildOperator(node planning.Node, timeRange types.QueryTimeRange, materializer *planning.Materializer, side string) (types.InstantVectorOperator, types.ScalarOperator, error) {
+	o, err := materializer.ConvertNodeToOperator(node, timeRange)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	switch o := o.(type) {
 	case types.InstantVectorOperator:
-		return child, nil
+		return o, nil, nil
 	case types.ScalarOperator:
-		return nil, child
+		return nil, o, nil
 	default:
-		return nil, nil
+		return nil, nil, fmt.Errorf("expected either InstantVectorOperator or ScalarOperator on %s-hand side of BinaryExpression, got %T", side, o)
 	}
 }
 

pkg/streamingpromql/planning/core/binary_expression.go

@@ -82,12 +82,22 @@ func (f *FunctionCall) ChildrenLabels() []string {
 	return l
 }
 
-func (f *FunctionCall) OperatorFactory(children []types.Operator, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
+func (f *FunctionCall) OperatorFactory(materializer *planning.Materializer, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
 	fnc, ok := functions.RegisteredFunctions[f.Function]
 	if !ok {
 		return nil, compat.NewNotSupportedError(fmt.Sprintf("'%v' function", f.Function.PromQLName()))
 	}
 
+	children := make([]types.Operator, 0, len(f.Args))
+	for _, arg := range f.Args {
+		o, err := materializer.ConvertNodeToOperator(arg, timeRange)
+		if err != nil {
+			return nil, err
+		}
+
+		children = append(children, o)
+	}
+
 	var absentLabels labels.Labels
 
 	if f.Function == functions.FUNCTION_ABSENT || f.Function == functions.FUNCTION_ABSENT_OVER_TIME {

pkg/streamingpromql/planning/core/function_call.go

@@ -61,7 +61,7 @@ func (m *MatrixSelector) ChildrenLabels() []string {
 	return nil
 }
 
-func (m *MatrixSelector) OperatorFactory(_ []types.Operator, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
+func (m *MatrixSelector) OperatorFactory(_ *planning.Materializer, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
 	matchers, err := LabelMatchersToPrometheusType(m.Matchers)
 	if err != nil {
 		return nil, err

pkg/streamingpromql/planning/core/matrix_selector.go

@@ -56,7 +56,7 @@ func (n *NumberLiteral) ChildrenLabels() []string {
 	return nil
 }
 
-func (n *NumberLiteral) OperatorFactory(_ []types.Operator, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
+func (n *NumberLiteral) OperatorFactory(_ *planning.Materializer, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
 	o := scalars.NewScalarConstant(n.Value, timeRange, params.MemoryConsumptionTracker, n.ExpressionPosition.ToPrometheusType())
 
 	return planning.NewSingleUseOperatorFactory(o), nil

pkg/streamingpromql/planning/core/number_literal.go

@@ -56,7 +56,7 @@ func (s *StringLiteral) ChildrenLabels() []string {
 	return nil
 }
 
-func (s *StringLiteral) OperatorFactory(_ []types.Operator, _ types.QueryTimeRange, _ *planning.OperatorParameters) (planning.OperatorFactory, error) {
+func (s *StringLiteral) OperatorFactory(_ *planning.Materializer, _ types.QueryTimeRange, _ *planning.OperatorParameters) (planning.OperatorFactory, error) {
 	o := operators.NewStringLiteral(s.Value, s.ExpressionPosition.ToPrometheusType())
 
 	return planning.NewSingleUseOperatorFactory(o), nil

pkg/streamingpromql/planning/core/string_literal.go

@@ -48,11 +48,6 @@ func (s *Subquery) Describe() string {
 }
 
 func (s *Subquery) ChildrenTimeRange(timeRange types.QueryTimeRange) types.QueryTimeRange {
-	return SubqueryTimeRange(timeRange, s.Range, s.Step, s.Timestamp, s.Offset)
-}
-
-// FIXME: inline this into the method above once we no longer support directly converting from PromQL expressions to operators
-func SubqueryTimeRange(parentTimeRange types.QueryTimeRange, rng time.Duration, step time.Duration, ts *time.Time, offset time.Duration) types.QueryTimeRange {
 	// Subqueries are evaluated as a single range query with steps aligned to Unix epoch time 0.
 	// They are not evaluated as queries aligned to the individual step timestamps.
 	// See https://www.robustperception.io/promql-subqueries-and-alignment/ for an explanation.
@@ -67,24 +62,23 @@ func SubqueryTimeRange(parentTimeRange types.QueryTimeRange, rng time.Duration,
 	// This is relatively uncommon, and Prometheus' engine does the same thing. In the future, we
 	// could be smarter about this if it turns out to be a big problem.
 
-	start := parentTimeRange.StartT
-	end := parentTimeRange.EndT
-	stepMilliseconds := step.Milliseconds()
+	start := timeRange.StartT
+	end := timeRange.EndT
+	stepMilliseconds := s.Step.Milliseconds()
 
-	if ts != nil {
-		start = timestamp.FromTime(*ts)
+	if s.Timestamp != nil {
+		start = timestamp.FromTime(*s.Timestamp)
 		end = start
 	}
 
 	// Find the first timestamp inside the subquery range that is aligned to the step.
-	alignedStart := stepMilliseconds * ((start - offset.Milliseconds() - rng.Milliseconds()) / stepMilliseconds)
-	if alignedStart < start-offset.Milliseconds()-rng.Milliseconds() {
+	alignedStart := stepMilliseconds * ((start - s.Offset.Milliseconds() - s.Range.Milliseconds()) / stepMilliseconds)
+	if alignedStart < start-s.Offset.Milliseconds()-s.Range.Milliseconds() {
 		alignedStart += stepMilliseconds
 	}
 
-	end = end - offset.Milliseconds()
-
-	return types.NewRangeQueryTimeRange(timestamp.Time(alignedStart), timestamp.Time(end), step)
+	end = end - s.Offset.Milliseconds()
+	return types.NewRangeQueryTimeRange(timestamp.Time(alignedStart), timestamp.Time(end), s.Step)
 }
 
 func (s *Subquery) Details() proto.Message {
@@ -124,16 +118,14 @@ func (s *Subquery) ChildrenLabels() []string {
 	return []string{""}
 }
 
-func (s *Subquery) OperatorFactory(children []types.Operator, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
-	if len(children) != 1 {
-		return nil, fmt.Errorf("expected exactly 1 child for Subquery, got %v", len(children))
+func (s *Subquery) OperatorFactory(materializer *planning.Materializer, timeRange types.QueryTimeRange, params *planning.OperatorParameters) (planning.OperatorFactory, error) {
+	innerTimeRange := s.ChildrenTimeRange(timeRange)
+	inner, err := materializer.ConvertNodeToInstantVectorOperator(s.Inner, innerTimeRange)
+	if err != nil {
+		return nil, fmt.Errorf("could not create inner operator for Subquery: %w", err)
 	}
 
-	inner, ok := children[0].(types.InstantVectorOperator)
-	if !ok {
-		return nil, fmt.Errorf("expected InstantVectorOperator as child of Subquery, got %T", children[0])
-	}
-	o, err := operators.NewSubquery(inner, timeRange, s.ChildrenTimeRange(timeRange), TimestampFromTime(s.Timestamp), s.Offset, s.Range, s.ExpressionPosition.ToPrometheusType(), params.MemoryConsumptionTracker)
+	o, err := operators.NewSubquery(inner, timeRange, innerTimeRange, TimestampFromTime(s.Timestamp), s.Offset, s.Range, s.ExpressionPosition.ToPrometheusType(), params.MemoryConsumptionTracker)
 	if err != nil {
 		return nil, err
 	}