fix sorting in sort_by_label and sort_by_label_desc

CHANGELOG.md

@@ -9,6 +9,7 @@
 * [BUGFIX] Update prometheus alertmanager version to v0.34.1 to apply upstream inhibit rule bug fixes. #7877
 * [BUGFIX] Tenant Federation: Fix regex tenant federation not working in single binary mode, and failing with `multiple org IDs present` when a matched tenant ID contains regex metacharacters and also matches other tenants as a regex. #7869
 * [BUGFIX] Query Frontend: Fix the data selection time window computed for an instant query sent without a `time` parameter. The default evaluation time was passed to `util.ParseTimeParam` in milliseconds, but that helper formats the default and parses it as seconds, so the value overflowed `UnixNano()` and wrapped to a timestamp decades away from `now`. This made the `data_select_min_time` and `data_select_max_time` query stats wrong, and stopped the `time_window` query attribute of query priority and query rejection from matching such queries correctly. #7880
+* [BUGFIX] Query Frontend: Fix vertically sharded instant queries using `sort_by_label` or `sort_by_label_desc` returning series ordered by sample value instead of by the given label values. #7886
 
 ## 1.22.0 in progress
 * [CHANGE] Ruler: Remove the deprecated `-ruler.evaluation-delay-duration` flag and its `ruler_evaluation_delay_duration` per-tenant limit. Use `-ruler.query-offset` / `ruler_query_offset`, which no longer takes the higher of the two values. Cortex decodes the runtime config strictly, so a leftover `ruler_evaluation_delay_duration` override makes the runtime config fail to load: Cortex **exits at startup** (`module failed`, `module=runtime-config`), and on an already-running process every reload fails, pinning the last good overrides and dropping `cortex_runtime_config_last_reload_successful` to 0. Run `grep -r ruler_evaluation_delay_duration` over your runtime configs before upgrading. #7792

pkg/querier/tripperware/merge.go

@@ -7,6 +7,7 @@ import (
 	"sort"
 
 	"github.com/prometheus/common/model"
+	"github.com/prometheus/prometheus/promql"
 	promqlparser "github.com/prometheus/prometheus/promql/parser"
 	"github.com/thanos-io/thanos/pkg/strutil"
 
@@ -142,7 +143,7 @@ func matrixMerge(ctx context.Context, resps []*PrometheusResponse) ([]SampleStre
 func vectorMerge(ctx context.Context, req Request, resps []*PrometheusResponse) (*Vector, error) {
 	output := map[string]Sample{}
 	metrics := []string{} // Used to preserve the order for topk and bottomk.
-	sortPlan, err := sortPlanForQuery(req.GetQuery())
+	sortPlan, sortLabels, err := sortPlanForQuery(req.GetQuery())
 	if err != nil {
 		return nil, err
 	}
@@ -187,6 +188,13 @@ func vectorMerge(ctx context.Context, req Request, resps []*PrometheusResponse)
 		return result, nil
 	}
 
+	if sortPlan == sortByLabelValuesAsc || sortPlan == sortByLabelValuesDesc {
+		for _, metric := range sortByLabelValues(output, sortPlan, sortLabels) {
+			result.Samples = append(result.Samples, output[metric])
+		}
+		return result, nil
+	}
+
 	samples := make([]*pair, 0, len(output))
 	for k, v := range output {
 		samples = append(samples, &pair{
@@ -266,10 +274,12 @@ func statsMerge(shouldSumStats bool, resps []*PrometheusResponse) *PrometheusRes
 type sortPlan int
 
 const (
-	mergeOnly        sortPlan = 0
-	sortByValuesAsc  sortPlan = 1
-	sortByValuesDesc sortPlan = 2
-	sortByLabels     sortPlan = 3
+	mergeOnly             sortPlan = 0
+	sortByValuesAsc       sortPlan = 1
+	sortByValuesDesc      sortPlan = 2
+	sortByLabels          sortPlan = 3
+	sortByLabelValuesAsc  sortPlan = 4
+	sortByLabelValuesDesc sortPlan = 5
 )
 
 type pair struct {
@@ -288,60 +298,96 @@ func getSortValueFromPair(samples []*pair, i int) float64 {
 	return samples[i].s.Sample.Value
 }
 
-func sortPlanForQuery(q string) (sortPlan, error) {
+// sortByLabelValues returns the keys of output ordered by Prometheus' own
+// sort_by_label or sort_by_label_desc implementation.
+func sortByLabelValues(output map[string]Sample, plan sortPlan, sortLabels []string) []string {
+	funcName := "sort_by_label"
+	if plan == sortByLabelValuesDesc {
+		funcName = "sort_by_label_desc"
+	}
+	// The first argument is the vector, which is passed via vectorVals instead.
+	args := make(promqlparser.Expressions, 0, len(sortLabels)+1)
+	args = append(args, nil)
+	for _, l := range sortLabels {
+		args = append(args, &promqlparser.StringLiteral{Val: l})
+	}
+
+	vec := make(promql.Vector, 0, len(output))
+	for _, s := range output {
+		vec = append(vec, promql.Sample{Metric: cortexpb.FromLabelAdaptersToLabels(s.Labels)})
+	}
+	vec, _ = promql.FunctionCalls[funcName]([]promql.Vector{vec}, nil, args, &promql.EvalNodeHelper{})
+
+	buf := make([]byte, 0, 1024)
+	metrics := make([]string, 0, len(vec))
+	for _, s := range vec {
+		metrics = append(metrics, string(s.Metric.Bytes(buf)))
+	}
+	return metrics
+}
+
+// sortPlanForQuery returns how merged vector results should be ordered. For
+// sort_by_label and sort_by_label_desc, it also returns the labels to sort by.
+func sortPlanForQuery(q string) (sortPlan, []string, error) {
 	expr, err := cortexparser.ParseExpr(q)
 	if err != nil {
-		return 0, err
+		return 0, nil, err
 	}
 	// Check if the root expression is topk, bottomk, limitk, or limit_ratio
 	if aggr, ok := expr.(*promqlparser.AggregateExpr); ok {
 		if aggr.Op == promqlparser.TOPK || aggr.Op == promqlparser.BOTTOMK || aggr.Op == promqlparser.LIMITK || aggr.Op == promqlparser.LIMIT_RATIO {
-			return mergeOnly, nil
+			return mergeOnly, nil, nil
 		}
 	}
-	checkForSort := func(expr promqlparser.Expr) (sortAsc, sortDesc bool) {
-		if n, ok := expr.(*promqlparser.Call); ok {
-			if n.Func != nil {
-				if n.Func.Name == "sort" {
-					sortAsc = true
-				}
-				if n.Func.Name == "sort_desc" {
-					sortDesc = true
-				}
-				if n.Func.Name == "sort_by_label" {
-					sortAsc = true
-				}
-				if n.Func.Name == "sort_by_label_desc" {
-					sortDesc = true
-				}
-			}
+	checkForSort := func(expr promqlparser.Expr) (sortPlan, []string, bool) {
+		n, ok := expr.(*promqlparser.Call)
+		if !ok || n.Func == nil {
+			return 0, nil, false
+		}
+		switch n.Func.Name {
+		case "sort":
+			return sortByValuesAsc, nil, true
+		case "sort_desc":
+			return sortByValuesDesc, nil, true
+		case "sort_by_label":
+			return sortByLabelValuesAsc, stringsFromArgs(n.Args[1:]), true
+		case "sort_by_label_desc":
+			return sortByLabelValuesDesc, stringsFromArgs(n.Args[1:]), true
 		}
-		return sortAsc, sortDesc
+		return 0, nil, false
 	}
 	// Check the root expression for sort
-	if sortAsc, sortDesc := checkForSort(expr); sortAsc || sortDesc {
-		if sortAsc {
-			return sortByValuesAsc, nil
-		}
-		return sortByValuesDesc, nil
+	if plan, sortLabels, ok := checkForSort(expr); ok {
+		return plan, sortLabels, nil
 	}
 
 	// If the root expression is a binary expression, check the LHS and RHS for sort
 	if bin, ok := expr.(*promqlparser.BinaryExpr); ok {
-		if sortAsc, sortDesc := checkForSort(bin.LHS); sortAsc || sortDesc {
-			if sortAsc {
-				return sortByValuesAsc, nil
-			}
-			return sortByValuesDesc, nil
+		if plan, sortLabels, ok := checkForSort(bin.LHS); ok {
+			return plan, sortLabels, nil
+		}
+		if plan, sortLabels, ok := checkForSort(bin.RHS); ok {
+			return plan, sortLabels, nil
 		}
-		if sortAsc, sortDesc := checkForSort(bin.RHS); sortAsc || sortDesc {
-			if sortAsc {
-				return sortByValuesAsc, nil
+	}
+	return sortByLabels, nil, nil
+}
+
+func stringsFromArgs(args promqlparser.Expressions) []string {
+	strs := make([]string, 0, len(args))
+	for _, arg := range args {
+		for {
+			p, ok := arg.(*promqlparser.ParenExpr)
+			if !ok {
+				break
 			}
-			return sortByValuesDesc, nil
+			arg = p.Expr
+		}
+		if s, ok := arg.(*promqlparser.StringLiteral); ok {
+			strs = append(strs, s.Val)
 		}
 	}
-	return sortByLabels, nil
+	return strs
 }
 
 // mergeSampleStreams deduplicates sample streams using a map.

pkg/querier/tripperware/merge_test.go

@@ -1,11 +1,15 @@
 package tripperware
 
 import (
+	"context"
+	"fmt"
 	"testing"
 
+	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
 	promqlparser "github.com/prometheus/prometheus/promql/parser"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/cortexproject/cortex/pkg/cortexpb"
 	ingester_client "github.com/cortexproject/cortex/pkg/ingester/client"
@@ -597,9 +601,10 @@ func TestSliceHistograms(t *testing.T) {
 
 func Test_sortPlanForQuery(t *testing.T) {
 	tc := []struct {
-		query        string
-		expectedPlan sortPlan
-		err          bool
+		query          string
+		expectedPlan   sortPlan
+		expectedLabels []string
+		err            bool
 	}{
 		{
 			query:        "invalid(10, up)",
@@ -642,14 +647,22 @@ func Test_sortPlanForQuery(t *testing.T) {
 			err:          false,
 		},
 		{
-			query:        "1 + sort_by_label_desc(sum by (job) (up) )",
-			expectedPlan: sortByValuesDesc,
-			err:          false,
+			query:          `1 + sort_by_label_desc(sum by (job) (up), "job")`,
+			expectedPlan:   sortByLabelValuesDesc,
+			expectedLabels: []string{"job"},
+			err:            false,
 		},
 		{
-			query:        "sort_by_label(topk by (job) (10, up))",
-			expectedPlan: sortByValuesAsc,
-			err:          false,
+			query:          `sort_by_label(topk by (job) (10, up), "job", ("instance"))`,
+			expectedPlan:   sortByLabelValuesAsc,
+			expectedLabels: []string{"job", "instance"},
+			err:            false,
+		},
+		{
+			query:          "sort_by_label(sum by (job) (up))",
+			expectedPlan:   sortByLabelValuesAsc,
+			expectedLabels: []string{},
+			err:            false,
 		},
 		{
 			query:        "topk(5, up) by (job) + sort_desc(up)",
@@ -671,13 +684,115 @@ func Test_sortPlanForQuery(t *testing.T) {
 	for _, tc := range tc {
 		t.Run(tc.query, func(t *testing.T) {
 			promqlparser.EnableExperimentalFunctions = true
-			p, err := sortPlanForQuery(tc.query)
+			p, lbls, err := sortPlanForQuery(tc.query)
 			if tc.err {
 				assert.Error(t, err)
 			} else {
 				assert.NoError(t, err)
 				assert.Equal(t, tc.expectedPlan, p)
+				assert.Equal(t, tc.expectedLabels, lbls)
+			}
+		})
+	}
+}
+
+func TestMergeResponse_SortByLabel(t *testing.T) {
+	promqlparser.EnableExperimentalFunctions = true
+
+	sample := func(v float64, lbls ...string) Sample {
+		return Sample{
+			Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromStrings(lbls...)),
+			Sample: &cortexpb.Sample{Value: v, TimestampMs: 1000},
+		}
+	}
+	vectorResponse := func(samples ...Sample) Response {
+		return &PrometheusResponse{
+			Status: StatusSuccess,
+			Data: PrometheusData{
+				ResultType: model.ValVector.String(),
+				Result:     PrometheusQueryResult{Result: &PrometheusQueryResult_Vector{Vector: &Vector{Samples: samples}}},
+			},
+		}
+	}
+
+	for _, tc := range []struct {
+		name     string
+		query    string
+		resps    []Response
+		expected []Sample
+	}{
+		{
+			name:  "sort_by_label ignores sample values",
+			query: `sort_by_label(sum by (job) (up), "job")`,
+			resps: []Response{
+				vectorResponse(sample(4, "job", "a"), sample(1, "job", "c")),
+				vectorResponse(sample(3, "job", "b"), sample(2, "job", "d")),
+			},
+			expected: []Sample{sample(4, "job", "a"), sample(3, "job", "b"), sample(1, "job", "c"), sample(2, "job", "d")},
+		},
+		{
+			name:  "sort_by_label_desc ignores sample values",
+			query: `sort_by_label_desc(sum by (job) (up), "job")`,
+			resps: []Response{
+				vectorResponse(sample(1, "job", "c"), sample(4, "job", "a")),
+				vectorResponse(sample(2, "job", "d"), sample(3, "job", "b")),
+			},
+			expected: []Sample{sample(2, "job", "d"), sample(1, "job", "c"), sample(3, "job", "b"), sample(4, "job", "a")},
+		},
+		{
+			name:  "sort_by_label uses natural sort order",
+			query: `sort_by_label(sum by (job) (up), "job")`,
+			resps: []Response{
+				vectorResponse(sample(1, "job", "job10"), sample(2, "job", "job1")),
+				vectorResponse(sample(3, "job", "job2")),
+			},
+			expected: []Sample{sample(2, "job", "job1"), sample(3, "job", "job2"), sample(1, "job", "job10")},
+		},
+		{
+			name:  "sort_by_label sorts by the given labels before the full label set",
+			query: `sort_by_label(sum by (env, job) (up), "job", "env")`,
+			resps: []Response{
+				vectorResponse(sample(1, "env", "a", "job", "y"), sample(2, "env", "b", "job", "x")),
+				vectorResponse(sample(3, "env", "a", "job", "x"), sample(4, "env", "b", "job", "y")),
+			},
+			expected: []Sample{
+				sample(3, "env", "a", "job", "x"),
+				sample(2, "env", "b", "job", "x"),
+				sample(1, "env", "a", "job", "y"),
+				sample(4, "env", "b", "job", "y"),
+			},
+		},
+		{
+			name:  "sort_by_label falls back to the full label set when the given labels are equal",
+			query: `sort_by_label(sum by (env, job) (up), "env")`,
+			resps: []Response{
+				vectorResponse(sample(1, "env", "a", "job", "z")),
+				vectorResponse(sample(2, "env", "a", "job", "x")),
+			},
+			expected: []Sample{sample(2, "env", "a", "job", "x"), sample(1, "env", "a", "job", "z")},
+		},
+		{
+			name:  "sort_by_label in a binary expression",
+			query: `1 + sort_by_label(sum by (job) (up), "job")`,
+			resps: []Response{
+				vectorResponse(sample(3, "job", "b")),
+				vectorResponse(sample(4, "job", "a")),
+			},
+			expected: []Sample{sample(4, "job", "a"), sample(3, "job", "b")},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			resp, err := MergeResponse(context.Background(), true, &PrometheusRequest{Query: tc.query}, tc.resps...)
+			require.NoError(t, err)
+			// Compare label strings since reflect on LabelAdapter trips checkptr with stringlabels.
+			toStrings := func(samples []Sample) []string {
+				out := make([]string, 0, len(samples))
+				for _, s := range samples {
+					out = append(out, fmt.Sprintf("%s %v", cortexpb.FromLabelAdaptersToLabels(s.Labels).String(), s.Sample.Value))
+				}
+				return out
 			}
+			assert.Equal(t, toStrings(tc.expected), toStrings(resp.(*PrometheusResponse).Data.Result.GetVector().Samples))
 		})
 	}
 }