plan, engine: get rid of trim sorts optimizer

engine/engine.go

@@ -97,7 +97,7 @@ func (o Opts) getLogicalOptimizers() []logicalplan.Optimizer {
 		optimizers = make([]logicalplan.Optimizer, len(o.LogicalOptimizers))
 		copy(optimizers, o.LogicalOptimizers)
 	}
-	return append(optimizers, logicalplan.TrimSortFunctions{})
+	return optimizers
 }
 
 type remoteEngine struct {

logicalplan/distribute_test.go

@@ -128,22 +128,22 @@ max by (pod) (
 			name: "unsupported aggregation in the operand path",
 			expr: `max by (pod) (sort(quantile(0.9, http_requests_total)))`,
 			expected: `
-max by (pod) (sort(quantile(0.9,
+max by (pod) (quantile(0.9,
   dedup(
     remote(http_requests_total),
     remote(http_requests_total)
   )
-)))`,
+))`,
 		},
 		{
 			name: "binary operation in the operand path",
-			expr: `max by (pod) (sort(metric_a / metric_b))`,
+			expr: `max by (pod) (metric_a / metric_b)`,
 			expected: `
-max by (pod) (sort(
+max by (pod) (
   dedup(remote(metric_a), remote(metric_a)) 
   / 
   dedup(remote(metric_b), remote(metric_b))
-))`,
+)`,
 		},
 		{
 			name: "binary operation with aggregations",

logicalplan/plan.go

@@ -46,6 +46,9 @@ func New(expr parser.Expr, opts *query.Options) Plan {
 	setOffsetForAtModifier(opts.Start.UnixMilli(), expr)
 	setOffsetForInnerSubqueries(expr, opts)
 
+	// the engine handles sorting at the presentation layer
+	expr = trimSorts(expr)
+
 	return &plan{
 		expr: expr,
 		opts: opts,
@@ -143,6 +146,23 @@ func TraverseBottomUp(parent *parser.Expr, current *parser.Expr, transform func(
 	return true
 }
 
+func trimSorts(expr parser.Expr) parser.Expr {
+	TraverseBottomUp(nil, &expr, func(parent, current *parser.Expr) bool {
+		if current == nil || parent == nil {
+			return true
+		}
+		switch e := (*parent).(type) {
+		case *parser.Call:
+			switch e.Func.Name {
+			case "sort", "sort_desc":
+				*parent = *current
+			}
+		}
+		return false
+	})
+	return expr
+}
+
 // preprocessExpr wraps all possible step invariant parts of the given expression with
 // StepInvariantExpr. It also resolves the preprocessors.
 // Copied from Prometheus and adjusted to work with the vendored parser:

logicalplan/plan_test.go

@@ -138,6 +138,50 @@ func TestMatcherPropagation(t *testing.T) {
 	}
 }
 
+func TestTrimSorts(t *testing.T) {
+	cases := []struct {
+		name     string
+		expr     string
+		expected string
+	}{
+		// this test case is ok since the engine determines sorting order
+		// before running optimziers
+		{
+			name:     "simple sort",
+			expr:     "sort(X)",
+			expected: "X",
+		},
+		{
+			name:     "sort",
+			expr:     "sum(sort(X))",
+			expected: "sum(X)",
+		},
+		{
+			name:     "nested",
+			expr:     "sum(sort(rate(X[1m])))",
+			expected: "sum(rate(X[1m]))",
+		},
+		{
+			name:     "weirdly nested",
+			expr:     "sum(sort(sqrt(sort(X))))",
+			expected: "sum(sqrt(X))",
+		},
+		{
+			name:     "sort in binary expression",
+			expr:     "sort(sort(sqrt(X))/sort(sqrt(Y)))",
+			expected: "sqrt(X) / sqrt(Y)",
+		},
+	}
+	for _, tcase := range cases {
+		t.Run(tcase.name, func(t *testing.T) {
+			expr, err := parser.ParseExpr(tcase.expr)
+			testutil.Ok(t, err)
+
+			testutil.Equals(t, tcase.expected, trimSorts(expr).String())
+		})
+	}
+}
+
 func cleanUp(replacements map[string]*regexp.Regexp, expr string) string {
 	for replacement, match := range replacements {
 		expr = match.ReplaceAllString(expr, replacement)