Skip to content

Commit

Permalink
Fix sort for topk and bottomk
Browse files Browse the repository at this point in the history
Signed-off-by: 🌲 Harry 🌊 John 🏔 <johrry@amazon.com>
  • Loading branch information
harry671003 committed Feb 16, 2023
1 parent efab308 commit 6a53006
Show file tree
Hide file tree
Showing 2 changed files with 169 additions and 2 deletions.
11 changes: 11 additions & 0 deletions pkg/queryfrontend/queryinstant_codec.go
Original file line number Diff line number Diff line change
Expand Up @@ -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" {
Expand Down
160 changes: 158 additions & 2 deletions pkg/queryfrontend/queryinstant_codec_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 6a53006

Please sign in to comment.