Add fallback logic to thanos promql engine

Signed-off-by: SungJin1212 <[email protected]>

Author: SungJin1212

CHANGELOG.md

@@ -14,6 +14,7 @@
 * [ENHANCEMENT] Alertmanager: Add receiver validations for msteamsv2 and rocketchat. #6606
 * [ENHANCEMENT] Query Frontend: Add a `-frontend.enabled-ruler-query-stats` flag to configure whether to report the query stats log for queries coming from the Ruler. #6504
 * [ENHANCEMENT] OTLP: Support otlp metadata ingestion. #6617
+* [ENHANCEMENT] Querier: Add a fallback logic when use a thanos promql engine. #6630
 * [BUGFIX] Ingester: Avoid error or early throttling when READONLY ingesters are present in the ring #6517
 * [BUGFIX] Ingester: Fix labelset data race condition. #6573
 * [BUGFIX] Compactor: Cleaner should not put deletion marker for blocks with no-compact marker. #6576

go.mod

@@ -51,7 +51,7 @@ require (
 	github.com/spf13/afero v1.11.0
 	github.com/stretchr/testify v1.10.0
 	github.com/thanos-io/objstore v0.0.0-20241111205755-d1dd89d41f97
-	github.com/thanos-io/promql-engine v0.0.0-20250220213456-fab1185f8c6c
+	github.com/thanos-io/promql-engine v0.0.0-20250302135832-accbf0891a16
 	github.com/thanos-io/thanos v0.37.3-0.20250212101700-346d18bb0f80
 	github.com/uber/jaeger-client-go v2.30.0+incompatible
 	github.com/weaveworks/common v0.0.0-20230728070032-dd9e68f319d5

go.sum

@@ -1672,8 +1672,8 @@ github.com/thanos-community/galaxycache v0.0.0-20211122094458-3a32041a1f1e h1:f1
 github.com/thanos-community/galaxycache v0.0.0-20211122094458-3a32041a1f1e/go.mod h1:jXcofnrSln/cLI6/dhlBxPQZEEQHVPCcFaH75M+nSzM=
 github.com/thanos-io/objstore v0.0.0-20241111205755-d1dd89d41f97 h1:VjG0mwhN1DkncwDHFvrpd12/2TLfgYNRmEQA48ikp+0=
 github.com/thanos-io/objstore v0.0.0-20241111205755-d1dd89d41f97/go.mod h1:vyzFrBXgP+fGNG2FopEGWOO/zrIuoy7zt3LpLeezRsw=
-github.com/thanos-io/promql-engine v0.0.0-20250220213456-fab1185f8c6c h1:STCm5S4Aht3hOR0WQ0B3daZv21GQC13uPYIfkcN762U=
-github.com/thanos-io/promql-engine v0.0.0-20250220213456-fab1185f8c6c/go.mod h1:aHSV5hL94fNb7PklN9L0V10j+/RGIlzqbw7OLdNgZFs=
+github.com/thanos-io/promql-engine v0.0.0-20250302135832-accbf0891a16 h1:ezd8hNCWiGQr4kdfCHFa0VCSi+LAO/28Mna264nDs2c=
+github.com/thanos-io/promql-engine v0.0.0-20250302135832-accbf0891a16/go.mod h1:aHSV5hL94fNb7PklN9L0V10j+/RGIlzqbw7OLdNgZFs=
 github.com/thanos-io/thanos v0.37.3-0.20250212101700-346d18bb0f80 h1:mOCRYn9SLBWJCXAdP+qDfgZDc0eqDxDc2HZGKTZ5vzk=
 github.com/thanos-io/thanos v0.37.3-0.20250212101700-346d18bb0f80/go.mod h1:Y7D8la8B5rpzRVKq2HCR4hbYZ4LGroSPqIJjtizgQg8=
 github.com/tjhop/slog-gokit v0.1.2 h1:pmQI4SvU9h4gA0vIQsdhJQSqQg4mOmsPykG2/PM3j1I=

pkg/querier/engine_factory.go

@@ -0,0 +1,78 @@
+package querier
+
+import (
+	"context"
+	"time"
+
+	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/thanos-io/promql-engine/engine"
+	"github.com/thanos-io/promql-engine/logicalplan"
+)
+
+type EngineFactory struct {
+	prometheusEngine *promql.Engine
+	thanosEngine     *engine.Engine
+}
+
+func NewEngineFactory(opts promql.EngineOpts, enableThanosEngine bool) *EngineFactory {
+	prometheusEngine := promql.NewEngine(opts)
+
+	var thanosEngine *engine.Engine
+	if enableThanosEngine {
+		thanosEngine = engine.New(engine.Opts{
+			EngineOpts:        opts,
+			LogicalOptimizers: logicalplan.AllOptimizers,
+			EnableAnalysis:    true,
+		})
+	}
+
+	return &EngineFactory{
+		prometheusEngine: prometheusEngine,
+		thanosEngine:     thanosEngine,
+	}
+}
+
+func (qf *EngineFactory) NewInstantQuery(ctx context.Context, q storage.Queryable, opts promql.QueryOpts, qs string, ts time.Time) (promql.Query, error) {
+	if qf.thanosEngine != nil {
+		res, err := qf.thanosEngine.MakeInstantQuery(ctx, q, fromPromQLOpts(opts), qs, ts)
+		if err != nil {
+			if engine.IsUnimplemented(err) {
+				// fallback to use prometheus engine
+				goto fallback
+			}
+			return nil, err
+		}
+		return res, nil
+	}
+
+fallback:
+	return qf.prometheusEngine.NewInstantQuery(ctx, q, opts, qs, ts)
+}
+
+func (qf *EngineFactory) NewRangeQuery(ctx context.Context, q storage.Queryable, opts promql.QueryOpts, qs string, start, end time.Time, interval time.Duration) (promql.Query, error) {
+	if qf.thanosEngine != nil {
+		res, err := qf.thanosEngine.MakeRangeQuery(ctx, q, fromPromQLOpts(opts), qs, start, end, interval)
+		if err != nil {
+			if engine.IsUnimplemented(err) {
+				// fallback to use prometheus engine
+				goto fallback
+			}
+			return nil, err
+		}
+		return res, nil
+	}
+
+fallback:
+	return qf.prometheusEngine.NewRangeQuery(ctx, q, opts, qs, start, end, interval)
+}
+
+func fromPromQLOpts(opts promql.QueryOpts) *engine.QueryOpts {
+	if opts == nil {
+		return &engine.QueryOpts{}
+	}
+	return &engine.QueryOpts{
+		LookbackDeltaParam:      opts.LookbackDelta(),
+		EnablePerStepStatsParam: opts.EnablePerStepStats(),
+	}
+}

pkg/querier/querier.go

@@ -19,8 +19,6 @@ import (
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/util/annotations"
-	"github.com/thanos-io/promql-engine/engine"
-	"github.com/thanos-io/promql-engine/logicalplan"
 	"github.com/thanos-io/thanos/pkg/strutil"
 	"golang.org/x/sync/errgroup"
 
@@ -208,7 +206,6 @@ func New(cfg Config, limits *validation.Overrides, distributor Distributor, stor
 	// The cortex supports holt_winters for users using this function.
 	EnableExperimentalPromQLFunctions(cfg.EnablePromQLExperimentalFunctions, true)
 
-	var queryEngine promql.QueryEngine
 	opts := promql.EngineOpts{
 		Logger:               util_log.GoKitLogToSlog(logger),
 		Reg:                  reg,
@@ -223,15 +220,7 @@ func New(cfg Config, limits *validation.Overrides, distributor Distributor, stor
 			return cfg.DefaultEvaluationInterval.Milliseconds()
 		},
 	}
-	if cfg.ThanosEngine {
-		queryEngine = engine.New(engine.Opts{
-			EngineOpts:        opts,
-			LogicalOptimizers: logicalplan.AllOptimizers,
-			EnableAnalysis:    true,
-		})
-	} else {
-		queryEngine = promql.NewEngine(opts)
-	}
+	queryEngine := NewEngineFactory(opts, cfg.ThanosEngine)
 	return NewSampleAndChunkQueryable(lazyQueryable), exemplarQueryable, queryEngine
 }