engine: add sort_by_label functions

Signed-off-by: Michael Hoffmann <[email protected]>

engine/engine_test.go

@@ -40,6 +40,7 @@ import (
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/model/timestamp"
 	"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"
@@ -50,10 +51,15 @@ import (
 )
 
 func TestMain(m *testing.M) {
+	parser.EnableExperimentalFunctions = true
 	goleak.VerifyTestMain(m)
 }
 
 func TestPromqlAcceptance(t *testing.T) {
+	// promql acceptance tests disable experimental functions again
+	// since we use them in our tests too we need to enable them afterwards again
+	t.Cleanup(func() { parser.EnableExperimentalFunctions = true })
+
 	engine := engine.New(engine.Opts{
 		EngineOpts: promql.EngineOpts{
 			EnableAtModifier:         true,
@@ -2886,6 +2892,7 @@ func createVectorResult(vector promql.Vector) *promql.Result {
 
 func TestInstantQuery(t *testing.T) {
 	t.Parallel()
+
 	defaultQueryTime := time.Unix(50, 0)
 	cases := []struct {
 		load      string
@@ -4025,6 +4032,21 @@ min without () (
 >
   absent_over_time({__name__="http_requests_total",route="/"}[3m] offset 1m45s)`,
 		},
+		{
+			name: "sort_by_label",
+			load: `load 30s
+			    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`,
+			query: `sort_by_label_desc(http_requests, "instance")`,
+		},
 	}
 
 	disableOptimizerOpts := []bool{true, false}
@@ -4081,7 +4103,6 @@ min without () (
 								defer q2.Close()
 
 								oldResult := q2.Exec(ctx)
-
 								testutil.WithGoCmp(comparer).Equals(t, oldResult, newResult, queryExplanation(q1))
 							})
 						}

engine/sort.go

@@ -6,6 +6,7 @@ package engine
 import (
 	"math"
 
+	"github.com/facette/natsort"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/promql/parser"
@@ -26,6 +27,12 @@ type sortFuncResultSort struct {
 	sortOrder sortOrder
 }
 
+type sortByLabelFuncResult struct {
+	sortingLabels []string
+
+	sortOrder sortOrder
+}
+
 type aggregateResultSort struct {
 	sortingLabels []string
 	groupBy       bool
@@ -36,6 +43,16 @@ type aggregateResultSort struct {
 type noSortResultSort struct {
 }
 
+func extractSortingLabels(f *parser.Call) []string {
+	args := f.Args[1:]
+
+	res := make([]string, 0)
+	for i := range args {
+		res = append(res, args[i].(*parser.StringLiteral).Val)
+	}
+	return res
+}
+
 func newResultSort(expr parser.Expr) resultSorter {
 	switch texpr := expr.(type) {
 	case *parser.Call:
@@ -44,6 +61,10 @@ func newResultSort(expr parser.Expr) resultSorter {
 			return sortFuncResultSort{sortOrder: sortOrderAsc}
 		case "sort_desc":
 			return sortFuncResultSort{sortOrder: sortOrderDesc}
+		case "sort_by_label":
+			return sortByLabelFuncResult{sortOrder: sortOrderAsc, sortingLabels: extractSortingLabels(texpr)}
+		case "sort_by_label_desc":
+			return sortByLabelFuncResult{sortOrder: sortOrderDesc, sortingLabels: extractSortingLabels(texpr)}
 		}
 	case *parser.AggregateExpr:
 		switch texpr.Op {
@@ -84,6 +105,28 @@ func (s sortFuncResultSort) comparer(samples *promql.Vector) func(i, j int) bool
 	}
 }
 
+func (s sortByLabelFuncResult) comparer(samples *promql.Vector) func(i, j int) bool {
+	return func(i, j int) bool {
+		iLb := labels.NewBuilder((*samples)[i].Metric)
+		jLb := labels.NewBuilder((*samples)[j].Metric)
+
+		for _, label := range s.sortingLabels {
+			lv1 := iLb.Get(label)
+			lv2 := jLb.Get(label)
+
+			if lv1 == lv2 {
+				continue
+			}
+			if natsort.Compare(lv1, lv2) {
+				return s.sortOrder == sortOrderAsc
+			} else {
+				return s.sortOrder == sortOrderDesc
+			}
+		}
+		return valueCompare(s.sortOrder, (*samples)[i].F, (*samples)[j].F)
+	}
+}
+
 func (s aggregateResultSort) comparer(samples *promql.Vector) func(i, j int) bool {
 	return func(i int, j int) bool {
 		var iLbls labels.Labels

execution/function/functions.go

@@ -161,6 +161,12 @@ var instantVectorFuncs = map[string]functionCall{
 	"sort_desc": simpleFunc(func(v float64) float64 {
 		return v
 	}),
+	"sort_by_label": simpleFunc(func(v float64) float64 {
+		return v
+	}),
+	"sort_by_label_desc": simpleFunc(func(v float64) float64 {
+		return v
+	}),
 }
 
 type noArgFunctionCall func(t int64) float64

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/cespare/xxhash/v2 v2.3.0
 	github.com/cortexproject/promqlsmith v0.0.0-20240506042652-6cfdd9739a5e
 	github.com/efficientgo/core v1.0.0-rc.2
+	github.com/facette/natsort v0.0.0-20181210072756-2cd4dd1e2dcb
 	github.com/go-kit/log v0.2.1
 	github.com/google/go-cmp v0.6.0
 	github.com/prometheus/client_golang v1.19.1
@@ -33,7 +34,6 @@ require (
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/dennwc/varint v1.0.0 // indirect
 	github.com/edsrzf/mmap-go v1.1.0 // indirect
-	github.com/facette/natsort v0.0.0-20181210072756-2cd4dd1e2dcb // indirect
 	github.com/go-logfmt/logfmt v0.6.0 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect

logicalplan/plan.go

@@ -269,7 +269,7 @@ func trimSorts(expr Node) Node {
 		switch e := (*parent).(type) {
 		case *FunctionCall:
 			switch e.Func.Name {
-			case "sort", "sort_desc":
+			case "sort", "sort_desc", "sort_by_label", "sort_by_label_desc":
 				*parent = *current
 			}
 		}