try disable binary expression for not on and no labels

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

Author: yeya24

pkg/cortex/modules.go

@@ -511,7 +511,7 @@ func (t *Cortex) initFlusher() (serv services.Service, err error) {
 // initQueryFrontendTripperware instantiates the tripperware used by the query frontend
 // to optimize Prometheus query requests.
 func (t *Cortex) initQueryFrontendTripperware() (serv services.Service, err error) {
-	queryAnalyzer := tripperware.NewDisableWithoutNameAnalyzer(querysharding.NewQueryAnalyzer())
+	queryAnalyzer := tripperware.NewDisableBinaryExpressionAnalyzer(querysharding.NewQueryAnalyzer())
 	// PrometheusCodec is a codec to encode and decode Prometheus query range requests and responses.
 	prometheusCodec := queryrange.NewPrometheusCodec(false, t.Cfg.Querier.ResponseCompression, t.Cfg.API.QuerierDefaultCodec)
 	// ShardedPrometheusCodec is same as PrometheusCodec but to be used on the sharded queries (it sum up the stats)

pkg/querier/tripperware/shard_by.go

@@ -3,22 +3,27 @@ package tripperware
 import (
 	"context"
 	"net/http"
-	"slices"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/prometheus/prometheus/model/labels"
+	"github.com/pkg/errors"
+	promqlparser "github.com/prometheus/prometheus/promql/parser"
 	"github.com/thanos-io/thanos/pkg/querysharding"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
 	"github.com/weaveworks/common/httpgrpc"
 
+	"github.com/cortexproject/cortex/pkg/parser"
 	querier_stats "github.com/cortexproject/cortex/pkg/querier/stats"
 	cquerysharding "github.com/cortexproject/cortex/pkg/querysharding"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
+var (
+	stop = errors.New("stop")
+)
+
 func ShardByMiddleware(logger log.Logger, limits Limits, merger Merger, queryAnalyzer querysharding.Analyzer) Middleware {
 	return MiddlewareFunc(func(next Handler) Handler {
 		return shardBy{
@@ -127,22 +132,33 @@ func InjectVerticalShardSizeToContext(ctx context.Context, verticalShardSize int
 	return context.WithValue(ctx, verticalShardsKey{}, verticalShardSize)
 }
 
-type disableWithoutNameAnalyzer struct {
+type disableBinaryExpressionAnalyzer struct {
 	analyzer querysharding.Analyzer
 }
 
-func NewDisableWithoutNameAnalyzer(analyzer querysharding.Analyzer) *disableWithoutNameAnalyzer {
-	return &disableWithoutNameAnalyzer{analyzer: analyzer}
+func NewDisableBinaryExpressionAnalyzer(analyzer querysharding.Analyzer) *disableBinaryExpressionAnalyzer {
+	return &disableBinaryExpressionAnalyzer{analyzer: analyzer}
 }
 
-func (d *disableWithoutNameAnalyzer) Analyze(query string) (querysharding.QueryAnalysis, error) {
+func (d *disableBinaryExpressionAnalyzer) Analyze(query string) (querysharding.QueryAnalysis, error) {
 	analysis, err := d.analyzer.Analyze(query)
-	if err != nil || !analysis.IsShardable() || analysis.ShardBy() {
+	if err != nil || !analysis.IsShardable() {
 		return analysis, err
 	}
 
-	// We are only interested in not shard by case.
-	if slices.Contains(analysis.ShardingLabels(), labels.MetricName) {
+	expr, _ := parser.ParseExpr(query)
+	isShardable := true
+	promqlparser.Inspect(expr, func(node promqlparser.Node, nodes []promqlparser.Node) error {
+		switch n := node.(type) {
+		case *promqlparser.BinaryExpr:
+			if !n.VectorMatching.On && len(n.VectorMatching.MatchingLabels) == 0 {
+				isShardable = false
+				return stop
+			}
+		}
+		return nil
+	})
+	if !isShardable {
 		// Mark as not shardable.
 		return querysharding.QueryAnalysis{}, nil
 	}