promql: add sort_by_label and sort_by_label_desc functions

This adds functions to sort a vector by its label value.

Based on https://github.com/prometheus/prometheus/pull/1533

Signed-off-by: Alexander Trost <galexrt@googlemail.com>

cmd/prometheus/main.go

@@ -154,6 +154,7 @@ type flagConfig struct {
 	enableNewSDManager         bool
 	enablePerStepStats         bool
 	enableAutoGOMAXPROCS       bool
+	enablePromQLSortByLabel    bool
 
 	prometheusURL   string
 	corsRegexString string
@@ -209,6 +210,8 @@ func (c *flagConfig) setFeatureListOptions(logger log.Logger) error {
 				config.DefaultConfig.GlobalConfig.ScrapeProtocols = config.DefaultNativeHistogramScrapeProtocols
 				config.DefaultGlobalConfig.ScrapeProtocols = config.DefaultNativeHistogramScrapeProtocols
 				level.Info(logger).Log("msg", "Experimental native histogram support enabled. Changed default scrape_protocols to prefer PrometheusProto format.", "global.scrape_protocols", fmt.Sprintf("%v", config.DefaultGlobalConfig.ScrapeProtocols))
+			case "promql-sort-by-label":
+				c.enablePromQLSortByLabel = true
 			case "":
 				continue
 			case "promql-at-modifier", "promql-negative-offset":
@@ -665,6 +668,7 @@ func main() {
 			EnableAtModifier:     true,
 			EnableNegativeOffset: true,
 			EnablePerStepStats:   cfg.enablePerStepStats,
+			EnableSortByLabel:    cfg.enablePromQLSortByLabel,
 		}
 
 		queryEngine = promql.NewEngine(opts)

docs/feature_flags.md

@@ -195,3 +195,10 @@ won't work when you push OTLP metrics.
 
 Enables PromQL functions that are considered experimental and whose name or
 semantics could change.
+
+## PromQL: Sort By Label Function
+
+`--enable-feature=promql-sort-by-label`
+
+When enabled, the `sort_by_label` and `sort_by_label_desc` functions can be used
+to sort instant query results by their label values.

docs/querying/functions.md

@@ -586,6 +586,22 @@ in ascending order. Native histograms are sorted by their sum of observations.
 
 Same as `sort`, but sorts in descending order.
 
+## `sort_by_label()`
+
+**This function has to be enabled via a [feature flag](../feature_flags/#promql-sort-by-label-function).**
+
+`sort_by_label(v instant-vector, label string, ...)` returns vector elements sorted by their label values and sample value in case of label values being equal, in ascending order.
+
+Please note that the sort by label functions only affect the results of instant queries, as range query results always have a fixed output ordering.
+
+## `sort_by_label_desc()`
+
+**This function has to be enabled via a [feature flag](../feature_flags/#promql-sort-by-label-function).**
+
+Same as `sort_by_label`, but sorts in descending order.
+
+Please note that the sort by label functions only affect the results of instant queries, as range query results always have a fixed output ordering.
+
 ## `sqrt()`
 
 `sqrt(v instant-vector)` calculates the square root of all elements in `v`.

promql/engine.go

@@ -304,6 +304,9 @@ type EngineOpts struct {
 
 	// EnablePerStepStats if true allows for per-step stats to be computed on request. Disabled otherwise.
 	EnablePerStepStats bool
+
+	// EnableSortByLabel if true enables sort_by_label/sort_by_label_desc query functions.
+	EnableSortByLabel bool
 }
 
 // Engine handles the lifetime of queries from beginning to end.
@@ -321,6 +324,7 @@ type Engine struct {
 	enableAtModifier         bool
 	enableNegativeOffset     bool
 	enablePerStepStats       bool
+	enableSortByLabel        bool
 }
 
 // NewEngine returns a new engine.
@@ -411,6 +415,7 @@ func NewEngine(opts EngineOpts) *Engine {
 		enableAtModifier:         opts.EnableAtModifier,
 		enableNegativeOffset:     opts.EnableNegativeOffset,
 		enablePerStepStats:       opts.EnablePerStepStats,
+		enableSortByLabel:        opts.EnableSortByLabel,
 	}
 }
 
@@ -702,6 +707,7 @@ func (ng *Engine) execEvalStmt(ctx context.Context, query *query, s *parser.Eval
 			lookbackDelta:            s.LookbackDelta,
 			samplesStats:             query.sampleStats,
 			noStepSubqueryIntervalFn: ng.noStepSubqueryIntervalFn,
+			enableSortByLabel:        ng.enableSortByLabel,
 		}
 		query.sampleStats.InitStepTracking(start, start, 1)
 
@@ -759,6 +765,7 @@ func (ng *Engine) execEvalStmt(ctx context.Context, query *query, s *parser.Eval
 		lookbackDelta:            s.LookbackDelta,
 		samplesStats:             query.sampleStats,
 		noStepSubqueryIntervalFn: ng.noStepSubqueryIntervalFn,
+		enableSortByLabel:        ng.enableSortByLabel,
 	}
 	query.sampleStats.InitStepTracking(evaluator.startTimestamp, evaluator.endTimestamp, evaluator.interval)
 	val, warnings, err := evaluator.Eval(s.Expr)
@@ -1012,6 +1019,7 @@ type evaluator struct {
 	lookbackDelta            time.Duration
 	samplesStats             *stats.QuerySamples
 	noStepSubqueryIntervalFn func(rangeMillis int64) int64
+	enableSortByLabel        bool
 }
 
 // errorf causes a panic with the input formatted into an error.
@@ -1387,6 +1395,8 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 			if ok {
 				return ev.rangeEvalTimestampFunctionOverVectorSelector(vs, call, e)
 			}
+		} else if (e.Func.Name == "sort_by_label" || e.Func.Name == "sort_by_label_desc") && !ev.enableSortByLabel {
+			ev.error(fmt.Errorf("sort_by_label() and sort_by_label_desc() require the \"promql-sort-by-label\" feature flag to be set"))
 		}
 
 		// Check if the function has a matrix argument.
@@ -1780,6 +1790,7 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 			lookbackDelta:            ev.lookbackDelta,
 			samplesStats:             ev.samplesStats.NewChild(),
 			noStepSubqueryIntervalFn: ev.noStepSubqueryIntervalFn,
+			enableSortByLabel:        ev.enableSortByLabel,
 		}
 		res, ws := newEv.eval(e.Expr)
 		ev.currentSamples = newEv.currentSamples

promql/functions.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/grafana/regexp"
 	"github.com/prometheus/common/model"
+	"golang.org/x/exp/slices"
 
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
@@ -367,6 +368,68 @@ func funcSortDesc(vals []parser.Value, args parser.Expressions, enh *EvalNodeHel
 	return Vector(byValueSorter), nil
 }
 
+// === sort_by_label(vector parser.ValueTypeVector, label parser.ValueTypeString...) (Vector, Annotations) ===
+func funcSortByLabel(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelper) (Vector, annotations.Annotations) {
+	// In case the labels are the same, NaN should sort to the bottom, so take
+	// ascending sort with NaN first and reverse it.
+	var anno annotations.Annotations
+	vals[0], anno = funcSort(vals, args, enh)
+	labels := stringSliceFromArgs(args[1:])
+	slices.SortFunc(vals[0].(Vector), func(a, b Sample) int {
+		// Iterate over each given label
+		for _, label := range labels {
+			lv1 := a.Metric.Get(label)
+			lv2 := b.Metric.Get(label)
+			// 0 if a == b, -1 if a < b, and +1 if a > b.
+			switch strings.Compare(lv1, lv2) {
+			case -1:
+				return -1
+			case +1:
+				return +1
+			default:
+				continue
+			}
+		}
+
+		return 0
+	})
+
+	return vals[0].(Vector), anno
+}
+
+// === sort_by_label_desc(vector parser.ValueTypeVector, label parser.ValueTypeString...) (Vector, Annotations) ===
+func funcSortByLabelDesc(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelper) (Vector, annotations.Annotations) {
+	// In case the labels are the same, NaN should sort to the bottom, so take
+	// ascending sort with NaN first and reverse it.
+	var anno annotations.Annotations
+	vals[0], anno = funcSortDesc(vals, args, enh)
+	labels := stringSliceFromArgs(args[1:])
+	slices.SortFunc(vals[0].(Vector), func(a, b Sample) int {
+		// Iterate over each given label
+		for _, label := range labels {
+			lv1 := a.Metric.Get(label)
+			lv2 := b.Metric.Get(label)
+			// If label values are the same, continue to the next label
+			if lv1 == lv2 {
+				continue
+			}
+			// 0 if a == b, -1 if a < b, and +1 if a > b.
+			switch strings.Compare(lv1, lv2) {
+			case -1:
+				return +1
+			case +1:
+				return -1
+			default:
+				continue
+			}
+		}
+
+		return 0
+	})
+
+	return vals[0].(Vector), anno
+}
+
 // === clamp(Vector parser.ValueTypeVector, min, max Scalar) (Vector, Annotations) ===
 func funcClamp(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelper) (Vector, annotations.Annotations) {
 	vec := vals[0].(Vector)
@@ -1493,6 +1556,8 @@ var FunctionCalls = map[string]FunctionCall{
 	"sinh":               funcSinh,
 	"sort":               funcSort,
 	"sort_desc":          funcSortDesc,
+	"sort_by_label":      funcSortByLabel,
+	"sort_by_label_desc": funcSortByLabelDesc,
 	"sqrt":               funcSqrt,
 	"stddev_over_time":   funcStddevOverTime,
 	"stdvar_over_time":   funcStdvarOverTime,
@@ -1637,3 +1702,11 @@ func stringFromArg(e parser.Expr) string {
 	unwrapParenExpr(&tmp)             // Optionally unwrap ParenExpr
 	return tmp.(*parser.StringLiteral).Val
 }
+
+func stringSliceFromArgs(args parser.Expressions) []string {
+	tmp := make([]string, len(args))
+	for i := 0; i < len(args); i++ {
+		tmp[i] = stringFromArg(args[i])
+	}
+	return tmp
+}

promql/parser/functions.go

@@ -347,6 +347,18 @@ var Functions = map[string]*Function{
 		ArgTypes:   []ValueType{ValueTypeVector},
 		ReturnType: ValueTypeVector,
 	},
+	"sort_by_label": {
+		Name:       "sort_by_label",
+		ArgTypes:   []ValueType{ValueTypeVector, ValueTypeString},
+		Variadic:   -1,
+		ReturnType: ValueTypeVector,
+	},
+	"sort_by_label_desc": {
+		Name:       "sort_by_label_desc",
+		ArgTypes:   []ValueType{ValueTypeVector, ValueTypeString},
+		Variadic:   -1,
+		ReturnType: ValueTypeVector,
+	},
 	"sqrt": {
 		Name:       "sqrt",
 		ArgTypes:   []ValueType{ValueTypeVector},

promql/promql_test.go

@@ -35,6 +35,7 @@ func newTestEngine() *Engine {
 		EnableAtModifier:         true,
 		EnableNegativeOffset:     true,
 		EnablePerStepStats:       true,
+		EnableSortByLabel:        true,
 	})
 }
 

promql/testdata/functions.test

@@ -469,6 +469,116 @@ eval_ordered instant at 50m sort_desc(http_requests)
 	http_requests{group="production", instance="0", job="api-server"} 100
 	http_requests{group="canary", instance="2", job="api-server"} NaN
 
+# Tests for sort_by_label/sort_by_label_desc.
+clear
+load 5m
+	http_requests{job="api-server", instance="0", group="production"}	0+10x10
+	http_requests{job="api-server", instance="1", group="production"}	0+20x10
+	http_requests{job="api-server", instance="0", group="canary"}		0+30x10
+	http_requests{job="api-server", instance="1", group="canary"}		0+40x10
+	http_requests{job="api-server", instance="2", group="canary"}		NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN
+	http_requests{job="app-server", instance="0", group="production"}	0+50x10
+	http_requests{job="app-server", instance="1", group="production"}	0+60x10
+	http_requests{job="app-server", instance="0", group="canary"}		0+70x10
+	http_requests{job="app-server", instance="1", group="canary"}		0+80x10
+	http_requests{job="api-server", instance="2", group="production"}	0+10x10
+
+eval_ordered instant at 50m sort_by_label(http_requests, "instance")
+	http_requests{group="production", instance="0", job="api-server"} 100
+	http_requests{group="canary", instance="0", job="api-server"} 300
+	http_requests{group="production", instance="0", job="app-server"} 500
+	http_requests{group="canary", instance="0", job="app-server"} 700
+	http_requests{group="production", instance="1", job="api-server"} 200
+	http_requests{group="canary", instance="1", job="api-server"} 400
+	http_requests{group="production", instance="1", job="app-server"} 600
+	http_requests{group="canary", instance="1", job="app-server"} 800
+	http_requests{group="production", instance="2", job="api-server"} 100
+	http_requests{group="canary", instance="2", job="api-server"} NaN
+
+eval_ordered instant at 50m sort_by_label(http_requests, "instance", "group")
+	http_requests{group="canary", instance="0", job="api-server"} 300
+	http_requests{group="canary", instance="0", job="app-server"} 700
+	http_requests{group="production", instance="0", job="api-server"} 100
+	http_requests{group="production", instance="0", job="app-server"} 500
+	http_requests{group="canary", instance="1", job="api-server"} 400
+	http_requests{group="canary", instance="1", job="app-server"} 800
+	http_requests{group="production", instance="1", job="api-server"} 200
+	http_requests{group="production", instance="1", job="app-server"} 600
+	http_requests{group="canary", instance="2", job="api-server"} NaN
+	http_requests{group="production", instance="2", job="api-server"} 100
+
+eval_ordered instant at 50m sort_by_label(http_requests, "instance", "group")
+	http_requests{group="canary", instance="0", job="api-server"} 300
+	http_requests{group="canary", instance="0", job="app-server"} 700
+	http_requests{group="production", instance="0", job="api-server"} 100
+	http_requests{group="production", instance="0", job="app-server"} 500
+	http_requests{group="canary", instance="1", job="api-server"} 400
+	http_requests{group="canary", instance="1", job="app-server"} 800
+	http_requests{group="production", instance="1", job="api-server"} 200
+	http_requests{group="production", instance="1", job="app-server"} 600
+	http_requests{group="canary", instance="2", job="api-server"} NaN
+	http_requests{group="production", instance="2", job="api-server"} 100
+
+eval_ordered instant at 50m sort_by_label(http_requests, "group", "instance", "job")
+	http_requests{group="canary", instance="0", job="api-server"} 300
+	http_requests{group="canary", instance="0", job="app-server"} 700
+	http_requests{group="canary", instance="1", job="api-server"} 400
+	http_requests{group="canary", instance="1", job="app-server"} 800
+	http_requests{group="canary", instance="2", job="api-server"} NaN
+	http_requests{group="production", instance="0", job="api-server"} 100
+	http_requests{group="production", instance="0", job="app-server"} 500
+	http_requests{group="production", instance="1", job="api-server"} 200
+	http_requests{group="production", instance="1", job="app-server"} 600
+	http_requests{group="production", instance="2", job="api-server"} 100
+
+eval_ordered instant at 50m sort_by_label(http_requests, "job", "instance", "group")
+	http_requests{group="canary", instance="0", job="api-server"} 300
+	http_requests{group="production", instance="0", job="api-server"} 100
+	http_requests{group="canary", instance="1", job="api-server"} 400
+	http_requests{group="production", instance="1", job="api-server"} 200
+	http_requests{group="canary", instance="2", job="api-server"} NaN
+	http_requests{group="production", instance="2", job="api-server"} 100
+	http_requests{group="canary", instance="0", job="app-server"} 700
+	http_requests{group="production", instance="0", job="app-server"} 500
+	http_requests{group="canary", instance="1", job="app-server"} 800
+	http_requests{group="production", instance="1", job="app-server"} 600
+
+eval_ordered instant at 50m sort_by_label_desc(http_requests, "instance")
+	http_requests{group="production", instance="2", job="api-server"} 100
+	http_requests{group="canary", instance="2", job="api-server"} NaN
+	http_requests{group="canary", instance="1", job="app-server"} 800
+	http_requests{group="production", instance="1", job="app-server"} 600
+	http_requests{group="canary", instance="1", job="api-server"} 400
+	http_requests{group="production", instance="1", job="api-server"} 200
+	http_requests{group="canary", instance="0", job="app-server"} 700
+	http_requests{group="production", instance="0", job="app-server"} 500
+	http_requests{group="canary", instance="0", job="api-server"} 300
+	http_requests{group="production", instance="0", job="api-server"} 100
+
+eval_ordered instant at 50m sort_by_label_desc(http_requests, "instance", "group")
+	http_requests{group="production", instance="2", job="api-server"} 100
+	http_requests{group="canary", instance="2", job="api-server"} NaN
+	http_requests{group="production", instance="1", job="app-server"} 600
+	http_requests{group="production", instance="1", job="api-server"} 200
+	http_requests{group="canary", instance="1", job="app-server"} 800
+	http_requests{group="canary", instance="1", job="api-server"} 400
+	http_requests{group="production", instance="0", job="app-server"} 500
+	http_requests{group="production", instance="0", job="api-server"} 100
+	http_requests{group="canary", instance="0", job="app-server"} 700
+	http_requests{group="canary", instance="0", job="api-server"} 300
+
+eval_ordered instant at 50m sort_by_label_desc(http_requests, "instance", "group", "job")
+	http_requests{group="production", instance="2", job="api-server"} 100
+	http_requests{group="canary", instance="2", job="api-server"} NaN
+	http_requests{group="production", instance="1", job="app-server"} 600
+	http_requests{group="production", instance="1", job="api-server"} 200
+	http_requests{group="canary", instance="1", job="app-server"} 800
+	http_requests{group="canary", instance="1", job="api-server"} 400
+	http_requests{group="production", instance="0", job="app-server"} 500
+	http_requests{group="production", instance="0", job="api-server"} 100
+	http_requests{group="canary", instance="0", job="app-server"} 700
+	http_requests{group="canary", instance="0", job="api-server"} 300
+
 # Tests for holt_winters
 clear
 

web/ui/module/codemirror-promql/src/complete/promql.terms.ts

@@ -427,6 +427,18 @@ export const functionIdentifierTerms = [
     info: 'Sort input series descendingly by value',
     type: 'function',
   },
+  {
+    label: 'sort_by_label',
+    detail: 'function',
+    info: 'Sort input series ascendingly by label value',
+    type: 'function',
+  },
+  {
+    label: 'sort_by_label_desc',
+    detail: 'function',
+    info: 'Sort input series descendingly by value value',
+    type: 'function',
+  },
   {
     label: 'sqrt',
     detail: 'function',

web/ui/module/codemirror-promql/src/types/function.ts

@@ -74,6 +74,8 @@ import {
   Sinh,
   Sort,
   SortDesc,
+  SortByLabel,
+  SortByLabelDesc,
   Sqrt,
   StddevOverTime,
   StdvarOverTime,
@@ -476,6 +478,18 @@ const promqlFunctions: { [key: number]: PromQLFunction } = {
     variadic: 0,
     returnType: ValueType.vector,
   },
+  [SortByLabel]: {
+    name: 'sort_by_label',
+    argTypes: [ValueType.vector, ValueType.string],
+    variadic: -1,
+    returnType: ValueType.vector,
+  },
+  [SortByLabelDesc]: {
+    name: 'sort_by_label_desc',
+    argTypes: [ValueType.vector, ValueType.string],
+    variadic: -1,
+    returnType: ValueType.vector,
+  },
   [Sqrt]: {
     name: 'sqrt',
     argTypes: [ValueType.vector],

web/ui/module/lezer-promql/src/highlight.js

@@ -20,7 +20,7 @@ export const promQLHighLight = styleTags({
     NumberLiteral: tags.number,
     Duration: tags.number,
     Identifier: tags.variableName,
-    'Abs Absent AbsentOverTime Acos Acosh Asin Asinh Atan Atanh AvgOverTime Ceil Changes Clamp ClampMax ClampMin Cos Cosh CountOverTime DaysInMonth DayOfMonth DayOfWeek DayOfYear Deg Delta Deriv Exp Floor HistogramCount HistogramFraction HistogramQuantile HistogramSum HoltWinters Hour Idelta Increase Irate LabelReplace LabelJoin LastOverTime Ln Log10 Log2 MaxOverTime MinOverTime Minute Month Pi PredictLinear PresentOverTime QuantileOverTime Rad Rate Resets Round Scalar Sgn Sin Sinh Sort SortDesc Sqrt StddevOverTime StdvarOverTime SumOverTime Tan Tanh Time Timestamp Vector Year':
+    'Abs Absent AbsentOverTime Acos Acosh Asin Asinh Atan Atanh AvgOverTime Ceil Changes Clamp ClampMax ClampMin Cos Cosh CountOverTime DaysInMonth DayOfMonth DayOfWeek DayOfYear Deg Delta Deriv Exp Floor HistogramCount HistogramFraction HistogramQuantile HistogramSum HoltWinters Hour Idelta Increase Irate LabelReplace LabelJoin LastOverTime Ln Log10 Log2 MaxOverTime MinOverTime Minute Month Pi PredictLinear PresentOverTime QuantileOverTime Rad Rate Resets Round Scalar Sgn Sin Sinh Sort SortDesc SortByLabel SortByLabelDesc Sqrt StddevOverTime StdvarOverTime SumOverTime Tan Tanh Time Timestamp Vector Year':
         tags.function(tags.variableName),
     'Avg Bottomk Count Count_values Group Max Min Quantile Stddev Stdvar Sum Topk': tags.operatorKeyword,
     'By Without Bool On Ignoring GroupLeft GroupRight Offset Start End': tags.modifier,

web/ui/module/lezer-promql/src/promql.grammar

@@ -167,6 +167,8 @@ FunctionIdentifier {
   Sinh |
   Sort |
   SortDesc |
+  SortByLabel |
+  SortByLabelDesc |
   Sqrt |
   StddevOverTime |
   StdvarOverTime |
@@ -396,6 +398,8 @@ NumberLiteral  {
   Sinh { condFn<"sinh"> }
   Sort { condFn<"sort"> }
   SortDesc { condFn<"sort_desc"> }
+  SortByLabel { condFn<"sort_by_label"> }
+  SortByLabelDesc { condFn<"sort_by_label_desc"> }
   Sqrt { condFn<"sqrt"> }
   StddevOverTime { condFn<"stddev_over_time"> }
   StdvarOverTime { condFn<"stdvar_over_time"> }