First iteration of streaming PromQL engine

CHANGELOG.md

@@ -9,6 +9,7 @@
 * [FEATURE] Continuous-test: now runable as a module with `mimir -target=continuous-test`. #7747
 * [FEATURE] Store-gateway: Allow specific tenants to be enabled or disabled via `-store-gateway.enabled-tenants` or `-store-gateway.disabled-tenants` CLI flags or their corresponding YAML settings. #7653
 * [FEATURE] New `-<prefix>.s3.bucket-lookup-type` flag configures lookup style type, used to access bucket in s3 compatible providers. #7684
+* [FEATURE] Querier: add experimental streaming PromQL engine, enabled with `-querier.promql-engine=streaming`. #7693
 * [FEATURE] Server: added experimental [PROXY protocol support](https://www.haproxy.org/download/2.3/doc/proxy-protocol.txt). The PROXY protocol support can be enabled via `-server.proxy-protocol-enabled=true`. When enabled, the support is added both to HTTP and gRPC listening ports. #7698
 * [ENHANCEMENT] Store-gateway: merge series from different blocks concurrently. #7456
 * [ENHANCEMENT] Store-gateway: Add `stage="wait_max_concurrent"` to `cortex_bucket_store_series_request_stage_duration_seconds` which records how long the query had to wait for its turn for `-blocks-storage.bucket-store.max-concurrent`. #7609

Makefile

@@ -346,7 +346,7 @@ lint: check-makefiles
 	faillint -paths "github.com/grafana/mimir/pkg/..." ./pkg/alertmanager/alertspb/...
 	faillint -paths "github.com/grafana/mimir/pkg/..." ./pkg/ruler/rulespb/...
 	faillint -paths "github.com/grafana/mimir/pkg/..." ./pkg/storage/sharding/...
-	faillint -paths "github.com/grafana/mimir/pkg/..." ./pkg/querier/engine/...
+	faillint -paths "github.com/grafana/mimir/pkg/..." ./pkg/querier/engine/common/...
 	faillint -paths "github.com/grafana/mimir/pkg/..." ./pkg/querier/api/...
 	faillint -paths "github.com/grafana/mimir/pkg/..." ./pkg/util/globalerror
 

cmd/mimir/config-descriptor.json

@@ -1876,6 +1876,17 @@
           "fieldType": "duration",
           "fieldCategory": "advanced"
         },
+        {
+          "kind": "field",
+          "name": "promql_engine",
+          "required": false,
+          "desc": "PromQL engine to use, either 'standard' or 'streaming'",
+          "fieldValue": null,
+          "fieldDefaultValue": "standard",
+          "fieldFlag": "querier.promql-engine",
+          "fieldType": "string",
+          "fieldCategory": "experimental"
+        },
         {
           "kind": "field",
           "name": "max_concurrent",

cmd/mimir/help-all.txt.tmpl

@@ -1725,6 +1725,8 @@ Usage of ./cmd/mimir/mimir:
     	Delay before initiating requests to further ingesters when request minimization is enabled and the initially selected set of ingesters have not all responded. Ignored if -querier.minimize-ingester-requests is not enabled. (default 3s)
   -querier.prefer-streaming-chunks-from-store-gateways
     	[experimental] Request store-gateways stream chunks. Store-gateways will only respond with a stream of chunks if the target store-gateway supports this, and this preference will be ignored by store-gateways that do not support this.
+  -querier.promql-engine string
+    	[experimental] PromQL engine to use, either 'standard' or 'streaming' (default "standard")
   -querier.promql-experimental-functions-enabled
     	[experimental] Enable experimental PromQL functions. This config option should be set on query-frontend too when query sharding is enabled.
   -querier.query-ingesters-within duration

docs/sources/mimir/configure/about-versioning.md

@@ -126,6 +126,7 @@ The following features are currently experimental:
   - Maximum response size for active series queries (`-querier.active-series-results-max-size-bytes`)
   - Enable PromQL experimental functions (`-querier.promql-experimental-functions-enabled`)
   - Allow streaming of `/active_series` responses to the frontend (`-querier.response-streaming-enabled`)
+  - Streaming PromQL engine (`-querier.promql-engine=streaming`)
 - Query-frontend
   - `-query-frontend.querier-forget-delay`
   - Instant query splitting (`-query-frontend.split-instant-queries-by-interval`)

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

@@ -1325,6 +1325,10 @@ store_gateway_client:
 # CLI flag: -querier.minimize-ingester-requests-hedging-delay
 [minimize_ingester_requests_hedging_delay: <duration> | default = 3s]
 
+# (experimental) PromQL engine to use, either 'standard' or 'streaming'
+# CLI flag: -querier.promql-engine
+[promql_engine: <string> | default = "standard"]
+
 # The number of workers running in each querier process. This setting limits the
 # maximum number of concurrent queries in each querier.
 # CLI flag: -querier.max-concurrent

integration/querier_test.go

@@ -514,6 +514,52 @@ func TestQuerierWithBlocksStorageRunningInSingleBinaryMode(t *testing.T) {
 	}
 }
 
+func TestStreamingPromQLEngine(t *testing.T) {
+	s, err := e2e.NewScenario(networkName)
+	require.NoError(t, err)
+	defer s.Close()
+
+	flags := mergeFlags(BlocksStorageFlags(), BlocksStorageS3Flags(), map[string]string{
+		"-querier.promql-engine": "streaming",
+	})
+
+	consul := e2edb.NewConsul()
+	minio := e2edb.NewMinio(9000, flags["-blocks-storage.s3.bucket-name"])
+	require.NoError(t, s.StartAndWaitReady(consul, minio))
+
+	distributor := e2emimir.NewDistributor("distributor", consul.NetworkHTTPEndpoint(), flags)
+	ingester := e2emimir.NewIngester("ingester", consul.NetworkHTTPEndpoint(), flags)
+	querier := e2emimir.NewQuerier("querier", consul.NetworkHTTPEndpoint(), flags)
+	require.NoError(t, s.StartAndWaitReady(distributor, ingester, querier))
+
+	// Wait until the distributor and querier have updated the ring.
+	// The distributor should have 512 tokens for the ingester ring and 1 for the distributor ring,
+	// and the querier should have 512 tokens for the ingester ring.
+	require.NoError(t, distributor.WaitSumMetrics(e2e.Equals(512+1), "cortex_ring_tokens_total"))
+	require.NoError(t, querier.WaitSumMetrics(e2e.Equals(512), "cortex_ring_tokens_total"))
+
+	// Push a series to Mimir.
+	writeClient, err := e2emimir.NewClient(distributor.HTTPEndpoint(), "", "", "", "user-1")
+	require.NoError(t, err)
+
+	seriesName := "series_1"
+	seriesTimestamp := time.Now()
+	series, expectedVector, _ := generateFloatSeries(seriesName, seriesTimestamp, prompb.Label{Name: seriesName, Value: seriesName})
+
+	res, err := writeClient.Push(series)
+	require.NoError(t, err)
+	require.Equal(t, 200, res.StatusCode)
+
+	// Query back the same series using the streaming PromQL engine.
+	c, err := e2emimir.NewClient("", querier.HTTPEndpoint(), "", "", "user-1")
+	require.NoError(t, err)
+
+	result, err := c.Query(seriesName, seriesTimestamp)
+	require.NoError(t, err)
+	require.Equal(t, model.ValVector, result.Type())
+	assert.Equal(t, expectedVector, result.(model.Vector))
+}
+
 func testMetadataQueriesWithBlocksStorage(
 	t *testing.T,
 	c *e2emimir.Client,

pkg/api/handlers.go

@@ -212,7 +212,7 @@ func NewQuerierHandler(
 	queryable storage.SampleAndChunkQueryable,
 	exemplarQueryable storage.ExemplarQueryable,
 	metadataSupplier querier.MetadataSupplier,
-	engine *promql.Engine,
+	engine promql.QueryEngine,
 	distributor Distributor,
 	reg prometheus.Registerer,
 	logger log.Logger,

pkg/mimir/mimir.go

@@ -709,7 +709,7 @@ type Mimir struct {
 	QuerierQueryable              prom_storage.SampleAndChunkQueryable
 	ExemplarQueryable             prom_storage.ExemplarQueryable
 	MetadataSupplier              querier.MetadataSupplier
-	QuerierEngine                 *promql.Engine
+	QuerierEngine                 promql.QueryEngine
 	QueryFrontendTripperware      querymiddleware.Tripperware
 	QueryFrontendCodec            querymiddleware.Codec
 	Ruler                         *ruler.Ruler

pkg/mimir/mimir_test.go

@@ -45,6 +45,7 @@ import (
 	"github.com/grafana/mimir/pkg/distributor"
 	"github.com/grafana/mimir/pkg/frontend/v1/frontendv1pb"
 	"github.com/grafana/mimir/pkg/ingester"
+	"github.com/grafana/mimir/pkg/querier"
 	"github.com/grafana/mimir/pkg/ruler"
 	"github.com/grafana/mimir/pkg/ruler/rulestore"
 	"github.com/grafana/mimir/pkg/scheduler/schedulerpb"
@@ -161,6 +162,9 @@ func TestMimir(t *testing.T) {
 			ReplicationFactor:      1,
 			InstanceInterfaceNames: []string{"en0", "eth0", "lo0", "lo"},
 		}},
+		Querier: querier.Config{
+			PromQLEngine: "standard",
+		},
 	}
 	require.NoError(t, cfg.Server.LogLevel.Set("info"))
 

pkg/mimir/modules.go

@@ -48,7 +48,7 @@ import (
 	"github.com/grafana/mimir/pkg/frontend/transport"
 	"github.com/grafana/mimir/pkg/ingester"
 	"github.com/grafana/mimir/pkg/querier"
-	"github.com/grafana/mimir/pkg/querier/engine"
+	"github.com/grafana/mimir/pkg/querier/engine/common"
 	"github.com/grafana/mimir/pkg/querier/tenantfederation"
 	querier_worker "github.com/grafana/mimir/pkg/querier/worker"
 	"github.com/grafana/mimir/pkg/ruler"
@@ -500,9 +500,12 @@ func (t *Mimir) initQueryable() (serv services.Service, err error) {
 	registerer := prometheus.WrapRegistererWith(querierEngine, t.Registerer)
 
 	// Create a querier queryable and PromQL engine
-	t.QuerierQueryable, t.ExemplarQueryable, t.QuerierEngine = querier.New(
+	t.QuerierQueryable, t.ExemplarQueryable, t.QuerierEngine, err = querier.New(
 		t.Cfg.Querier, t.Overrides, t.Distributor, t.StoreQueryable, registerer, util_log.Logger, t.ActivityTracker,
 	)
+	if err != nil {
+		return nil, fmt.Errorf("could not create queryable: %w", err)
+	}
 
 	// Use the distributor to return metric metadata by default
 	t.MetadataSupplier = t.Distributor
@@ -718,7 +721,7 @@ func (t *Mimir) initQueryFrontendCodec() (services.Service, error) {
 func (t *Mimir) initQueryFrontendTripperware() (serv services.Service, err error) {
 	promqlEngineRegisterer := prometheus.WrapRegistererWith(prometheus.Labels{"engine": "query-frontend"}, t.Registerer)
 
-	engineOpts, engineExperimentalFunctionsEnabled := engine.NewPromQLEngineOptions(t.Cfg.Querier.EngineConfig, t.ActivityTracker, util_log.Logger, promqlEngineRegisterer)
+	engineOpts, engineExperimentalFunctionsEnabled := common.NewPromQLEngineOptions(t.Cfg.Querier.EngineConfig, t.ActivityTracker, util_log.Logger, promqlEngineRegisterer)
 
 	tripperware, err := querymiddleware.NewTripperware(
 		t.Cfg.Frontend.QueryMiddleware,
@@ -842,7 +845,11 @@ func (t *Mimir) initRuler() (serv services.Service, err error) {
 		// TODO: Consider wrapping logger to differentiate from querier module logger
 		rulerRegisterer := prometheus.WrapRegistererWith(rulerEngine, t.Registerer)
 
-		queryable, _, eng := querier.New(t.Cfg.Querier, t.Overrides, t.Distributor, t.StoreQueryable, rulerRegisterer, util_log.Logger, t.ActivityTracker)
+		queryable, _, eng, err := querier.New(t.Cfg.Querier, t.Overrides, t.Distributor, t.StoreQueryable, rulerRegisterer, util_log.Logger, t.ActivityTracker)
+		if err != nil {
+			return nil, fmt.Errorf("could not create queryable for ruler: %w", err)
+		}
+
 		queryable = querier.NewErrorTranslateQueryableWithFn(queryable, ruler.WrapQueryableErrors)
 
 		if t.Cfg.Ruler.TenantFederation.Enabled {

pkg/querier/engine/config.go → pkg/querier/engine/common/config.go

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
-package engine
+package common
 
 import (
 	"flag"

pkg/querier/engine/query_tracker.go → pkg/querier/engine/common/query_tracker.go

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
-package engine
+package common
 
 import (
 	"bytes"

pkg/querier/engine/query_tracker_test.go → pkg/querier/engine/common/query_tracker_test.go

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
-package engine
+package common
 
 import (
 	"context"

pkg/querier/engine/streaming/README.md

@@ -0,0 +1,97 @@
+# Streaming PromQL engine
+
+This file contains a brief overview of the internals of the streaming PromQL engine.
+
+For an introduction to the engine itself and the problems it tries to solve, check out [this PromCon 2023 talk](https://www.youtube.com/watch?v=3kM2Asj6hcg).
+
+The goal of the engine is to allow evaluating queries over millions of series in a safe, performant and cost-effective way.
+To allow this, the engine aims to ensure that peak memory consumption of queriers is not proportional to the number of series selected.
+This will make it safe for operators to loosen the various query-related limits without risking the stability of their Mimir cluster or needing to devote enormous amounts of compute resources to queriers.
+
+The key way the engine achieves this is by not loading all the input series into memory at once, and instead streaming them into memory when needed.
+
+For example, let's say we're evaluating the query `sum by (environment) (some_metric{cluster="cluster-1"})`.
+
+Prometheus' PromQL engine will first load all samples for all series selected by `some_metric{cluster="cluster-1"}` into memory.
+It will then compute the sum for each unique value of `environment`.
+At its peak, Prometheus' PromQL engine will hold all samples for all input series (from `some_metric{cluster="cluster-1"}`) and all samples for all output series in memory at once.
+
+The streaming engine here will instead execute the selector `some_metric{cluster="cluster-1"}` and gather the labels of all series returned.
+With these labels, it will then compute all the possible output series for the `sum by (environment)` operation (ie. one output series per unique value of `environment`).
+Having computed the output series, it will then begin reading series from the selector, one at a time, and update the running total for the appropriate output series.
+At its peak, the streaming engine in this example will hold all samples for one input series and all samples for all output series in memory at once[^1],
+a significant reduction compared to Prometheus' PromQL engine, particularly when the selector selects many series.
+
+This idea of streaming can be applied to multiple levels as well. Imagine we're evaluating the query `max(sum by (environment) (some_metric{cluster="cluster-1"}))`.
+In the streaming engine, once the result of each group series produced by `sum` is complete, it is passed to `max`, which can update its running maximum seen so far across all groups.
+At its peak, the streaming engine will hold all samples for one input series, all samples for all incomplete `sum` group series, and the single incomplete `max` output series in memory at once.
+
+## Internals
+
+Within the streaming engine, a query is represented by a set of linked operators (one for each operation) that together form the query plan.
+
+For example, the `max(sum by (environment) (some_metric{cluster="cluster-1"}))` example from before would have a query plan made up of three operators:
+
+- The instant vector selector operator (`some_metric{cluster="cluster-1"}`)
+- The `sum` aggregation operator (`sum by (environment) (...)`), which consumes series from the instant vector selector operator
+- The `max` aggregation operator (`max (...)`), which consumes series from the `sum` aggregation operator
+
+Visually, the plan looks like this:
+
+```mermaid
+flowchart TB
+    IVS["`**instant vector selector**
+    some_metric#123;cluster=#quot;cluster-1#quot;#125;`"]
+    sum["`**sum aggregation**
+    sum by (environment) (...)`"]
+    max["`**max aggregation**
+    max (...)`"]
+    output((output))
+    IVS --> sum
+    sum --> max
+    max --> output
+```
+
+Each of these operators satisfies the `InstantVectorOperator` interface, defined [here](./operator/operator.go).
+The two key methods of this interface are `SeriesMetadata()` and `Next()`:
+
+`SeriesMetadata()` returns the list of all series' labels that will be returned by the operator[^2].
+In our example, the instant vector selector operator would return all the matching `some_metric` series, and the `sum` aggregation operator would return one series for each unique value of `environment`.
+
+`Next()` is then called by the consuming operator to read each series' data, one series at a time.
+In our example, the `sum` aggregation operator would call `Next()` on the instant vector selector operator to get the first series' data, then again to get the second series' data and so on.
+
+Elaborating on the example from before, the overall query would proceed like this, assuming the request is received over HTTP:
+
+1. query HTTP API handler calls `Engine.NewInstantQuery()` or `Engine.NewRangeQuery()` as appropriate ([source](./engine.go))
+   1. engine parses PromQL expression using Prometheus' PromQL parser, producing an abstract syntax tree (AST) ([source](./query.go))
+   1. engine converts AST produced by PromQL parser to query plan ([source](./query.go))
+   1. engine returns created `Query` instance
+1. query HTTP API handler calls `Query.Exec()`
+   1. `Query.Exec()` calls `SeriesMetadata()` on `max` aggregation operator
+      1. `max` aggregation operator calls `SeriesMetadata()` on `sum` aggregation operator
+         1. `sum` aggregation operator calls `SeriesMetadata()` on instant vector selector operator
+            - instant vector selector operator issues `Select()` call, which retrieves labels from ingesters and store-gateways
+         1. `sum` aggregation operator computes output series (one per unique value of `environment`) based on input series from instant vector selector
+      1. `max` aggregation operator computes output series based on input series from `sum` aggregation operator
+         - in this case, there's just one output series, given no grouping is being performed
+   1. root of the query calls `Next()` on `max` aggregation operator until all series have been returned
+      1. `max` aggregation operator calls `Next()` on `sum` aggregation operator
+         1. `sum` aggregation operator calls `Next()` on instant vector selector operator
+            - instant vector selector returns samples for next series
+         1. `sum` aggregation operator updates its running totals for the relevant output series
+         1. if all input series have now been seen for the output series just updated, `sum` aggregation operator returns that output series and removes it from its internal state
+         1. otherwise, it calls `Next()` again and repeats
+      1. `max` aggregation operator updates its running maximum based on the series returned
+      1. if all input series have been seen, `max` aggregation operator returns
+      1. otherwise, it calls `Next()` again and repeats
+1. query HTTP API handler converts returned result to wire format (either JSON or Protobuf) and sends to caller
+1. query HTTP API handler calls `Query.Close()` to release remaining resources
+
+[^1]:
+    This isn't strictly correct, as chunks streaming will buffer chunks for some series in memory as they're received over the network, and it ignores the initial memory consumption caused by the non-streaming calls to `SeriesMetadata()`.
+    But this applies equally to both engines when used in Mimir.
+
+[^2]:
+    This isn't done in a streaming fashion: all series' labels are loaded into memory at once.
+    In a future iteration of the engine, `SeriesMetadata()` could be made streaming as well, but this is out of scope for now.

pkg/querier/engine/streaming/compare.sh

@@ -0,0 +1,17 @@
+#! /usr/bin/env bash
+# SPDX-License-Identifier: AGPL-3.0-only
+# This script compares benchmark results for the two engines.
+
+set -euo pipefail
+
+RESULTS_FILE="$1" # Should be the path to a file produced by a command like `go test -run=XXX -bench="BenchmarkQuery" -count=6 -benchmem -timeout=1h .`
+
+STANDARD_RESULTS_FILE=$(mktemp /tmp/standard.XXXX)
+STREAMING_RESULTS_FILE=$(mktemp /tmp/streaming.XXXX)
+
+grep --invert-match "streaming-" "$RESULTS_FILE" | sed -E 's#/standard-[0-9]+##g' > "$STANDARD_RESULTS_FILE"
+grep --invert-match "standard-" "$RESULTS_FILE" | sed -E 's#/streaming-[0-9]+##g' > "$STREAMING_RESULTS_FILE"
+
+benchstat "$STANDARD_RESULTS_FILE" "$STREAMING_RESULTS_FILE" | sed "s#$STANDARD_RESULTS_FILE#     standard     #g" | sed "s#$STREAMING_RESULTS_FILE#     streaming     #g"
+
+rm "$STANDARD_RESULTS_FILE" "$STREAMING_RESULTS_FILE"