add projection pushdown optimizer

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

Author: yeya24

engine/pushdown_test.go

@@ -0,0 +1,274 @@
+package engine_test
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math/rand"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/cortexproject/promqlsmith"
+	"github.com/efficientgo/core/testutil"
+	"github.com/google/go-cmp/cmp"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/promql/parser"
+	"github.com/prometheus/prometheus/promql/promqltest"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/prometheus/prometheus/tsdb/chunkenc"
+	"github.com/prometheus/prometheus/util/annotations"
+	"github.com/thanos-io/promql-engine/engine"
+	"github.com/thanos-io/promql-engine/logicalplan"
+)
+
+type projectionQuerier struct {
+	storage.Querier
+}
+
+type projectionSeriesSet struct {
+	storage.SeriesSet
+	hints *storage.SelectHints
+}
+
+func (m projectionSeriesSet) Next() bool { return m.SeriesSet.Next() }
+func (m projectionSeriesSet) At() storage.Series {
+	// Get the original series
+	originalSeries := m.SeriesSet.At()
+	if originalSeries == nil {
+		return nil
+	}
+	// If no projection hints, return the original series
+	if m.hints == nil {
+		return originalSeries
+	}
+	if !m.hints.By && len(m.hints.Grouping) == 0 {
+		return originalSeries
+	}
+
+	// Apply projection based on hints
+	originalLabels := originalSeries.Labels()
+	var projectedLabels labels.Labels
+
+	if m.hints.By {
+		// Include mode: only keep the labels in the grouping
+		builder := labels.NewBuilder(labels.EmptyLabels())
+		for _, l := range originalLabels {
+			for _, groupLabel := range m.hints.Grouping {
+				if l.Name == groupLabel {
+					builder.Set(l.Name, l.Value)
+					break
+				}
+			}
+		}
+		builder.Set("__series_hash__", strconv.FormatUint(originalLabels.Hash(), 10))
+		projectedLabels = builder.Labels()
+	} else {
+		// Exclude mode: keep all labels except those in the grouping
+		excludeMap := make(map[string]struct{})
+		for _, groupLabel := range m.hints.Grouping {
+			excludeMap[groupLabel] = struct{}{}
+		}
+
+		builder := labels.NewBuilder(labels.EmptyLabels())
+		for _, l := range originalLabels {
+			if _, excluded := excludeMap[l.Name]; !excluded {
+				builder.Set(l.Name, l.Value)
+			}
+		}
+		builder.Set("__series_hash__", strconv.FormatUint(originalLabels.Hash(), 10))
+		projectedLabels = builder.Labels()
+	}
+
+	// Return a projected series that wraps the original but with filtered labels
+	return &projectedSeries{
+		Series: originalSeries,
+		lset:   projectedLabels,
+	}
+}
+
+// projectedSeries wraps a storage.Series but returns projected labels
+type projectedSeries struct {
+	storage.Series
+	lset labels.Labels
+}
+
+func (s *projectedSeries) Labels() labels.Labels {
+	return s.lset
+}
+
+func (s *projectedSeries) Iterator(iter chunkenc.Iterator) chunkenc.Iterator {
+	return s.Series.Iterator(iter)
+}
+
+func (m projectionSeriesSet) Err() error                        { return m.SeriesSet.Err() }
+func (m projectionSeriesSet) Warnings() annotations.Annotations { return m.SeriesSet.Warnings() }
+
+// Implement the Querier interface methods
+func (m *projectionQuerier) Select(ctx context.Context, sortSeries bool, hints *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
+	return projectionSeriesSet{
+		SeriesSet: m.Querier.Select(ctx, sortSeries, hints, matchers...),
+		hints:     hints,
+	}
+}
+func (m *projectionQuerier) LabelValues(ctx context.Context, name string, _ *storage.LabelHints, matchers ...*labels.Matcher) ([]string, annotations.Annotations, error) {
+	return nil, nil, nil
+}
+func (m *projectionQuerier) LabelNames(ctx context.Context, _ *storage.LabelHints, matchers ...*labels.Matcher) ([]string, annotations.Annotations, error) {
+	return nil, nil, nil
+}
+func (m *projectionQuerier) Close() error { return nil }
+
+// projectionQueryable is a storage.Queryable that applies projection to the querier
+type projectionQueryable struct {
+	storage.Queryable
+}
+
+func (q *projectionQueryable) Querier(mint, maxt int64) (storage.Querier, error) {
+	querier, err := q.Queryable.Querier(mint, maxt)
+	if err != nil {
+		return nil, err
+	}
+	return &projectionQuerier{
+		Querier: querier,
+	}, nil
+}
+
+func TestProjectionPushdownWithFuzz(t *testing.T) {
+	t.Parallel()
+
+	// Define test parameters
+	seed := time.Now().UnixNano()
+	rnd := rand.New(rand.NewSource(seed))
+	testRuns := 1000
+
+	// Create test data
+	load := `load 30s
+		http_requests_total{pod="nginx-1", job="app", env="prod", instance="1"} 1+1x40
+		http_requests_total{pod="nginx-2", job="app", env="dev", instance="2"} 2+2x40
+		http_requests_total{pod="nginx-3", job="api", env="prod", instance="3"} 3+3x40
+		http_requests_total{pod="nginx-4", job="api", env="dev", instance="4"} 4+4x40
+		errors_total{pod="nginx-1", job="app", env="prod", instance="1", cluster="us-west-2"} 0.5+0.5x40
+		errors_total{pod="nginx-2", job="app", env="dev", instance="2", cluster="us-west-2"} 1+1x40
+		errors_total{pod="nginx-3", job="api", env="prod", instance="3", cluster="us-east-2"} 1.5+1.5x40
+		errors_total{pod="nginx-4", job="api", env="dev", instance="4", cluster="us-east-1"} 2+2x40`
+
+	storage := promqltest.LoadedStorage(t, load)
+	defer storage.Close()
+
+	// Get series for PromQLSmith
+	seriesSet, err := getSeries(context.Background(), storage)
+	testutil.Ok(t, err)
+
+	// Configure PromQLSmith
+	psOpts := []promqlsmith.Option{
+		promqlsmith.WithEnableOffset(false),
+		promqlsmith.WithEnableAtModifier(false),
+		// Focus on aggregations that benefit from projection pushdown
+		promqlsmith.WithEnabledAggrs([]parser.ItemType{
+			parser.SUM, parser.MIN, parser.MAX, parser.AVG, parser.COUNT, parser.TOPK, parser.BOTTOMK,
+		}),
+		promqlsmith.WithEnableVectorMatching(true),
+	}
+	ps := promqlsmith.New(rnd, seriesSet, psOpts...)
+
+	// Engine options
+	engineOpts := promql.EngineOpts{
+		Timeout:              1 * time.Hour,
+		MaxSamples:           1e10,
+		EnableNegativeOffset: true,
+		EnableAtModifier:     true,
+	}
+
+	normalEngine := engine.New(engine.Opts{
+		EngineOpts:                  engineOpts,
+		LogicalOptimizers:           logicalplan.AllOptimizers,
+		DisableDuplicateLabelChecks: false,
+	})
+
+	pushdownEngine := engine.New(engine.Opts{
+		EngineOpts: engineOpts,
+		// ProjectionPushdown optimizer doesn't support merge selects optimizer
+		// so disable it for now.
+		LogicalOptimizers: []logicalplan.Optimizer{
+			logicalplan.SortMatchers{},
+			logicalplan.ProjectionPushdown{SeriesHashLabel: "__series_hash__"},
+			logicalplan.DetectHistogramStatsOptimizer{},
+			// logicalplan.MergeSelectsOptimizer{},
+		},
+		DisableDuplicateLabelChecks: false,
+	})
+
+	ctx := context.Background()
+	queryTime := time.Unix(600, 0)
+
+	t.Logf("Running %d fuzzy tests with seed %d", testRuns, seed)
+	for i := 0; i < testRuns; i++ {
+		var expr parser.Expr
+		var query string
+
+		// Generate a query that can be executed by the engine
+		for {
+			expr = ps.WalkInstantQuery()
+			query = expr.Pretty(0)
+
+			// Skip queries that don't benefit from projection pushdown
+			if !containsAggregationOrBinaryOperation(expr) {
+				continue
+			}
+
+			// Try to parse the query and see if it is valid.
+			_, err := normalEngine.NewInstantQuery(ctx, storage, nil, query, queryTime)
+			if err != nil {
+				continue
+			}
+			break
+		}
+
+		t.Run(fmt.Sprintf("Query_%d", i), func(t *testing.T) {
+			// Create projection querier that wraps the original querier
+			projectionStorage := &projectionQueryable{
+				Queryable: storage,
+			}
+
+			normalQuery, err := normalEngine.NewInstantQuery(ctx, storage, &engine.QueryOpts{}, query, queryTime)
+			testutil.Ok(t, err)
+			defer normalQuery.Close()
+			normalResult := normalQuery.Exec(ctx)
+			if normalResult.Err != nil {
+				// Something wrong with the generated query so it even failed without projection pushdown, skipping.
+				return
+			}
+			testutil.Ok(t, normalResult.Err, "query: %s", query)
+
+			pushdownQuery, err := pushdownEngine.MakeInstantQuery(ctx, projectionStorage, &engine.QueryOpts{}, query, queryTime)
+			testutil.Ok(t, err)
+
+			defer pushdownQuery.Close()
+			pushdownResult := pushdownQuery.Exec(ctx)
+			testutil.Ok(t, pushdownResult.Err, "query: %s", query)
+
+			if diff := cmp.Diff(normalResult, pushdownResult, comparer); diff != "" {
+				t.Errorf("Results differ for query %s: %s", query, diff)
+			}
+		})
+	}
+}
+
+// containsAggregationOrBinaryOperation checks if the expression contains any aggregation or binary operations
+func containsAggregationOrBinaryOperation(expr parser.Expr) bool {
+	found := false
+	parser.Inspect(expr, func(node parser.Node, _ []parser.Node) error {
+		switch node.(type) {
+		case *parser.AggregateExpr:
+			found = true
+			return errors.New("found")
+		case *parser.BinaryExpr:
+			found = true
+			return errors.New("found")
+		}
+		return nil
+	})
+	return found
+}

go.mod

@@ -1,10 +1,10 @@
 module github.com/thanos-io/promql-engine
 
-go 1.24
+go 1.24.0
 
 require (
 	github.com/cespare/xxhash/v2 v2.3.0
-	github.com/cortexproject/promqlsmith v0.0.0-20240506042652-6cfdd9739a5e
+	github.com/cortexproject/promqlsmith v0.0.0-20250407233056-90db95b1a4e4
 	github.com/efficientgo/core v1.0.0-rc.2
 	github.com/facette/natsort v0.0.0-20181210072756-2cd4dd1e2dcb
 	github.com/google/go-cmp v0.7.0

go.sum

@@ -38,8 +38,8 @@ github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UF
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cncf/xds/go v0.0.0-20240905190251-b4127c9b8d78 h1:QVw89YDxXxEe+l8gU8ETbOasdwEV+avkR75ZzsVV9WI=
 github.com/cncf/xds/go v0.0.0-20240905190251-b4127c9b8d78/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8=
-github.com/cortexproject/promqlsmith v0.0.0-20240506042652-6cfdd9739a5e h1:nOWmgQD3L/Z0bmm29iDxB7nlqjMnh7yD/PNOx9rnZmA=
-github.com/cortexproject/promqlsmith v0.0.0-20240506042652-6cfdd9739a5e/go.mod h1:+bSqRETXJ1uk2S93m//htzTVqu8DJPvlGEb3bSE9PzI=
+github.com/cortexproject/promqlsmith v0.0.0-20250407233056-90db95b1a4e4 h1:dpo7kQ24uFSV6Zgm9/kB34TIUWjGmadlbKrM6fNfQko=
+github.com/cortexproject/promqlsmith v0.0.0-20250407233056-90db95b1a4e4/go.mod h1:jh6POgN18lXU133HBMfwr/1TjvBp8e5kL4ZtRsAPvGY=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=

logicalplan/logical_nodes.go

@@ -88,6 +88,7 @@ func (f *VectorSelector) Clone() Node {
 	clone.Filters = shallowCloneSlice(f.Filters)
 	clone.LabelMatchers = shallowCloneSlice(f.LabelMatchers)
 	clone.Projection.Labels = shallowCloneSlice(f.Projection.Labels)
+	clone.Projection.Include = f.Projection.Include
 
 	if f.VectorSelector.Timestamp != nil {
 		ts := *f.VectorSelector.Timestamp

logicalplan/plan_test.go

@@ -6,6 +6,7 @@ package logicalplan
 import (
 	"fmt"
 	"regexp"
+	"sort"
 	"strings"
 	"testing"
 	"time"
@@ -39,6 +40,14 @@ func renderExprTree(expr Node) string {
 		if t.BatchSize > 0 {
 			base += fmt.Sprintf("[batch=%d]", t.BatchSize)
 		}
+		if t.Projection.Labels != nil {
+			sort.Strings(t.Projection.Labels)
+			if t.Projection.Include {
+				base += fmt.Sprintf("[projection=include(%s)]", strings.Join(t.Projection.Labels, ","))
+			} else {
+				base += fmt.Sprintf("[projection=exclude(%s)]", strings.Join(t.Projection.Labels, ","))
+			}
+		}
 		if len(t.Filters) > 0 {
 			b.WriteString("filter(")
 			b.WriteString(fmt.Sprintf("%s", t.Filters))
@@ -49,7 +58,10 @@ func renderExprTree(expr Node) string {
 		}
 		return base
 	case *MatrixSelector:
-		return t.String()
+		// Render the inner vector selector first
+		vsStr := renderExprTree(t.VectorSelector)
+		// Then add the range
+		return fmt.Sprintf("%s[%s]", vsStr, t.Range.String())
 	case *Binary:
 		var b strings.Builder
 		b.WriteString(renderExprTree(t.LHS))
@@ -108,6 +120,16 @@ func renderExprTree(expr Node) string {
 		return renderExprTree(t.Expr)
 	case *CheckDuplicateLabels:
 		return renderExprTree(t.Expr)
+	case *Subquery:
+		var b strings.Builder
+
+		// Render the inner expression
+		innerExpr := renderExprTree(t.Expr)
+		b.WriteString(innerExpr)
+
+		// Add the subquery range and step
+		b.WriteString(fmt.Sprintf("[%s:%s]", t.Range.String(), t.Step.String()))
+		return b.String()
 	default:
 		return t.String()
 	}