diff --git a/pkg/queryfrontend/queryinstant_codec.go b/pkg/queryfrontend/queryinstant_codec.go index 2f990a4899..b5adedebbd 100644 --- a/pkg/queryfrontend/queryinstant_codec.go +++ b/pkg/queryfrontend/queryinstant_codec.go @@ -338,6 +338,17 @@ func parseQueryForSort(q string) (bool, bool, error) { var sortDesc bool = false done := errors.New("done") promqlparser.Inspect(expr, func(n promqlparser.Node, _ []promqlparser.Node) error { + if n, ok := n.(*promqlparser.AggregateExpr); ok { + if n.Op == promqlparser.TOPK { + sortDesc = true + return done + } + if n.Op == promqlparser.BOTTOMK { + sortAsc = true + return done + } + return nil + } if n, ok := n.(*promqlparser.Call); ok { if n.Func != nil { if n.Func.Name == "sort" { diff --git a/pkg/queryfrontend/queryinstant_codec_test.go b/pkg/queryfrontend/queryinstant_codec_test.go index 0a86fa4f80..ac638a6da3 100644 --- a/pkg/queryfrontend/queryinstant_codec_test.go +++ b/pkg/queryfrontend/queryinstant_codec_test.go @@ -334,7 +334,7 @@ func TestMergeResponse(t *testing.T) { { name: "merge two responses with sort", req: &queryrange.PrometheusRequest{ - Query: "sort(up)", + Query: "1 + sort(topk(1, up))", }, resps: []queryrange.Response{ &queryrange.PrometheusInstantQueryResponse{ @@ -412,7 +412,7 @@ func TestMergeResponse(t *testing.T) { { name: "merge two responses with sort_desc", req: &queryrange.PrometheusRequest{ - Query: "sort_desc(up)", + Query: "1 + sort_desc(bottomk(1, up))", }, resps: []queryrange.Response{ &queryrange.PrometheusInstantQueryResponse{ @@ -487,6 +487,162 @@ func TestMergeResponse(t *testing.T) { }, }, }, + { + name: "merge two responses with topk", + req: &queryrange.PrometheusRequest{ + Query: "1 + topk(10, sort(up))", + }, + resps: []queryrange.Response{ + &queryrange.PrometheusInstantQueryResponse{ + Status: queryrange.StatusSuccess, + Data: queryrange.PrometheusInstantQueryData{ + ResultType: model.ValVector.String(), + Result: queryrange.PrometheusInstantQueryResult{ + Result: &queryrange.PrometheusInstantQueryResult_Vector{ + Vector: &queryrange.Vector{ + Samples: []*queryrange.Sample{ + { + Sample: cortexpb.Sample{TimestampMs: 0, Value: 1}, + Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{ + "__name__": "up", + "job": "foo", + })), + }, + }, + }, + }, + }, + }, + }, + &queryrange.PrometheusInstantQueryResponse{ + Status: queryrange.StatusSuccess, + Data: queryrange.PrometheusInstantQueryData{ + ResultType: model.ValVector.String(), + Result: queryrange.PrometheusInstantQueryResult{ + Result: &queryrange.PrometheusInstantQueryResult_Vector{ + Vector: &queryrange.Vector{ + Samples: []*queryrange.Sample{ + { + Sample: cortexpb.Sample{TimestampMs: 0, Value: 2}, + Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{ + "__name__": "up", + "job": "bar", + })), + }, + }, + }, + }, + }, + }, + }, + }, + expectedResp: &queryrange.PrometheusInstantQueryResponse{ + Status: queryrange.StatusSuccess, + Data: queryrange.PrometheusInstantQueryData{ + ResultType: model.ValVector.String(), + Result: queryrange.PrometheusInstantQueryResult{ + Result: &queryrange.PrometheusInstantQueryResult_Vector{ + Vector: &queryrange.Vector{ + Samples: []*queryrange.Sample{ + { + Sample: cortexpb.Sample{TimestampMs: 0, Value: 2}, + Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{ + "__name__": "up", + "job": "bar", + })), + }, + { + Sample: cortexpb.Sample{TimestampMs: 0, Value: 1}, + Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{ + "__name__": "up", + "job": "foo", + })), + }, + }, + }, + }, + }, + }, + }, + }, + { + name: "merge two responses with bottomk", + req: &queryrange.PrometheusRequest{ + Query: "1 + bottomk(10, sort(up))", + }, + resps: []queryrange.Response{ + &queryrange.PrometheusInstantQueryResponse{ + Status: queryrange.StatusSuccess, + Data: queryrange.PrometheusInstantQueryData{ + ResultType: model.ValVector.String(), + Result: queryrange.PrometheusInstantQueryResult{ + Result: &queryrange.PrometheusInstantQueryResult_Vector{ + Vector: &queryrange.Vector{ + Samples: []*queryrange.Sample{ + { + Sample: cortexpb.Sample{TimestampMs: 0, Value: 1}, + Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{ + "__name__": "up", + "job": "foo", + })), + }, + }, + }, + }, + }, + }, + }, + &queryrange.PrometheusInstantQueryResponse{ + Status: queryrange.StatusSuccess, + Data: queryrange.PrometheusInstantQueryData{ + ResultType: model.ValVector.String(), + Result: queryrange.PrometheusInstantQueryResult{ + Result: &queryrange.PrometheusInstantQueryResult_Vector{ + Vector: &queryrange.Vector{ + Samples: []*queryrange.Sample{ + { + Sample: cortexpb.Sample{TimestampMs: 0, Value: 2}, + Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{ + "__name__": "up", + "job": "bar", + })), + }, + }, + }, + }, + }, + }, + }, + }, + expectedResp: &queryrange.PrometheusInstantQueryResponse{ + Status: queryrange.StatusSuccess, + Data: queryrange.PrometheusInstantQueryData{ + ResultType: model.ValVector.String(), + Result: queryrange.PrometheusInstantQueryResult{ + Result: &queryrange.PrometheusInstantQueryResult_Vector{ + Vector: &queryrange.Vector{ + Samples: []*queryrange.Sample{ + { + Sample: cortexpb.Sample{TimestampMs: 0, Value: 1}, + Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{ + "__name__": "up", + "job": "foo", + })), + }, + { + Sample: cortexpb.Sample{TimestampMs: 0, Value: 2}, + Labels: cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{ + "__name__": "up", + "job": "bar", + })), + }, + }, + }, + }, + }, + }, + }, + }, { name: "merge two responses", req: defaultReq,