Fix sorted queries do not produce sorted results for shardable queries (#5148)

harry671003 · web-flow · commit 74d33df75ecd · 2023-02-15T12:59:46.000-08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -38,6 +38,7 @@
 * [BUGFIX] Ingester: Ingesters returning empty response for metadata APIs. #5081
 * [BUGFIX] Ingester: Fix panic when querying metadata from blocks that are being deleted. #5119
 * [BUGFIX] Ring: Fix case when dynamodb kv reaches the limit of 25 actions per batch call. #5136
+* [BUGFIX] Query-frontend:  Fix sorted queries do not produce sorted results for shardable queries. #5148
 * [FEATURE] Alertmanager: Add support for time_intervals. #5102
 
 ## 1.14.0 2022-12-02
diff --git a/pkg/querier/tripperware/instantquery/instant_query.go b/pkg/querier/tripperware/instantquery/instant_query.go
@@ -22,6 +22,8 @@ import (
 	"github.com/weaveworks/common/httpgrpc"
 	"google.golang.org/grpc/status"
 
+	promqlparser "github.com/prometheus/prometheus/promql/parser"
+
 	"github.com/cortexproject/cortex/pkg/cortexpb"
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 	"github.com/cortexproject/cortex/pkg/querier/tripperware/queryrange"
@@ -245,7 +247,7 @@ func (instantQueryCodec) EncodeResponse(ctx context.Context, res tripperware.Res
 	return &resp, nil
 }
 
-func (instantQueryCodec) MergeResponse(ctx context.Context, responses ...tripperware.Response) (tripperware.Response, error) {
+func (instantQueryCodec) MergeResponse(ctx context.Context, req tripperware.Request, responses ...tripperware.Response) (tripperware.Response, error) {
 	sp, _ := opentracing.StartSpanFromContext(ctx, "PrometheusInstantQueryResponse.MergeResponse")
 	sp.SetTag("response_count", len(responses))
 	defer sp.Finish()
@@ -265,11 +267,15 @@ func (instantQueryCodec) MergeResponse(ctx context.Context, responses ...tripper
 	// For now, we only shard queries that returns a vector.
 	switch promResponses[0].Data.ResultType {
 	case model.ValVector.String():
+		v, err := vectorMerge(req, promResponses)
+		if err != nil {
+			return nil, err
+		}
 		data = PrometheusInstantQueryData{
 			ResultType: model.ValVector.String(),
 			Result: PrometheusInstantQueryResult{
 				Result: &PrometheusInstantQueryResult_Vector{
-					Vector: vectorMerge(promResponses),
+					Vector: v,
 				},
 			},
 			Stats: statsMerge(promResponses),
@@ -297,8 +303,12 @@ func (instantQueryCodec) MergeResponse(ctx context.Context, responses ...tripper
 	return res, nil
 }
 
-func vectorMerge(resps []*PrometheusInstantQueryResponse) *Vector {
+func vectorMerge(req tripperware.Request, resps []*PrometheusInstantQueryResponse) (*Vector, error) {
 	output := map[string]*Sample{}
+	sortAsc, sortDesc, err := parseQueryForSort(req.GetQuery())
+	if err != nil {
+		return nil, err
+	}
 	buf := make([]byte, 0, 1024)
 	for _, resp := range resps {
 		if resp == nil {
@@ -327,22 +337,66 @@ func vectorMerge(resps []*PrometheusInstantQueryResponse) *Vector {
 	if len(output) == 0 {
 		return &Vector{
 			Samples: make([]*Sample, 0),
-		}
+		}, nil
 	}
 
-	keys := make([]string, 0, len(output))
-	for key := range output {
-		keys = append(keys, key)
+	type pair struct {
+		metric string
+		s      *Sample
+	}
+
+	samples := make([]*pair, 0, len(output))
+	for k, v := range output {
+		samples = append(samples, &pair{
+			metric: k,
+			s:      v,
+		})
 	}
-	sort.Strings(keys)
 
+	sort.Slice(samples, func(i, j int) bool {
+		// Order is determined by the sortFn in the query.
+		if sortAsc {
+			return samples[i].s.Sample.Value < samples[j].s.Sample.Value
+		} else if sortDesc {
+			return samples[i].s.Sample.Value > samples[j].s.Sample.Value
+		} else {
+			// Fallback on sorting by labels.
+			return samples[i].metric < samples[j].metric
+		}
+	})
 	result := &Vector{
 		Samples: make([]*Sample, 0, len(output)),
 	}
-	for _, key := range keys {
-		result.Samples = append(result.Samples, output[key])
+	for _, p := range samples {
+		result.Samples = append(result.Samples, p.s)
 	}
-	return result
+	return result, nil
+}
+
+func parseQueryForSort(q string) (bool, bool, error) {
+	expr, err := promqlparser.ParseExpr(q)
+	if err != nil {
+		return false, false, err
+	}
+	var sortAsc bool = false
+	var sortDesc bool = false
+	done := errors.New("done")
+	promqlparser.Inspect(expr, func(n promqlparser.Node, _ []promqlparser.Node) error {
+		if n, ok := n.(*promqlparser.Call); ok {
+			if n.Func != nil {
+				if n.Func.Name == "sort" {
+					sortAsc = true
+					return done
+				}
+				if n.Func.Name == "sort_desc" {
+					sortDesc = true
+					return done
+				}
+			}
+		}
+		return nil
+	})
+	return sortAsc, sortDesc, nil
 }
 
 func matrixMerge(resps []*PrometheusInstantQueryResponse) []tripperware.SampleStream {
diff --git a/pkg/querier/tripperware/instantquery/instant_query_test.go b/pkg/querier/tripperware/instantquery/instant_query_test.go
@@ -198,34 +198,43 @@ func TestResponse(t *testing.T) {
 }
 
 func TestMergeResponse(t *testing.T) {
+	defaultReq := &PrometheusRequest{
+		Query: "sum(up)",
+	}
 	for _, tc := range []struct {
 		name         string
+		req          tripperware.Request
 		resps        []string
 		expectedResp string
 		expectedErr  error
 	}{
 		{
 			name:         "empty response",
+			req:          defaultReq,
 			resps:        []string{`{"status":"success","data":{"resultType":"vector","result":[]}}`},
 			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[]}}`,
 		},
 		{
 			name:         "empty response with stats",
+			req:          defaultReq,
 			resps:        []string{`{"status":"success","data":{"resultType":"vector","result":[],"stats":{"samples":{"totalQueryableSamples":0,"totalQueryableSamplesPerStep":[]}}}}`},
 			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[],"stats":{"samples":{"totalQueryableSamples":0,"totalQueryableSamplesPerStep":[]}}}}`,
 		},
 		{
 			name:         "single response",
+			req:          defaultReq,
 			resps:        []string{`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}]}}`},
 			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}]}}`,
 		},
 		{
 			name:         "single response with stats",
+			req:          defaultReq,
 			resps:        []string{`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}],"stats":{"samples":{"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1,10]]}}}}`},
 			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}],"stats":{"samples":{"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1,10]]}}}}`,
 		},
 		{
 			name: "duplicated response",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}]}}`,
 				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}]}}`,
@@ -234,6 +243,7 @@ func TestMergeResponse(t *testing.T) {
 		},
 		{
 			name: "duplicated response with stats",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}],"stats":{"samples":{"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1,10]]}}}}`,
 				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}],"stats":{"samples":{"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1,10]]}}}}`,
@@ -242,14 +252,34 @@ func TestMergeResponse(t *testing.T) {
 		},
 		{
 			name: "merge two responses",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]}]}}`,
 				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"value":[2,"2"]}]}}`,
 			},
 			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"value":[2,"2"]},{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]}]}}`,
 		},
+		{
+			name: "merge two responses with sort",
+			req:  &PrometheusRequest{Query: "sort(up)"},
+			resps: []string{
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]}]}}`,
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"value":[1,"2"]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]},{"metric":{"__name__":"up","job":"bar"},"value":[1,"2"]}]}}`,
+		},
+		{
+			name: "merge two responses with sort_desc",
+			req:  &PrometheusRequest{Query: "sort_desc(up)"},
+			resps: []string{
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]}]}}`,
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"value":[1,"2"]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"value":[1,"2"]},{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]}]}}`,
+		},
 		{
 			name: "merge two responses with stats",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"value":[1,"1"]}],"stats":{"samples":{"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1,10]]}}}}`,
 				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"value":[2,"2"]}],"stats":{"samples":{"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1,10]]}}}}`,
@@ -258,6 +288,7 @@ func TestMergeResponse(t *testing.T) {
 		},
 		{
 			name: "responses don't contain vector, should return an error",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"string","result":[1662682521.409,"foo"]}}`,
 				`{"status":"success","data":{"resultType":"string","result":[1662682521.409,"foo"]}}`,
@@ -266,13 +297,15 @@ func TestMergeResponse(t *testing.T) {
 		},
 		{
 			name: "single matrix response",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"up"},"values":[[1,"1"],[2,"2"]]}]}}`,
 			},
 			expectedResp: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"up"},"values":[[1,"1"],[2,"2"]]}]}}`,
 		},
 		{
 			name: "multiple matrix responses without duplicated series",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"bar"},"values":[[1,"1"],[2,"2"]]}]}}`,
 				`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"foo"},"values":[[3,"3"],[4,"4"]]}]}}`,
@@ -281,6 +314,7 @@ func TestMergeResponse(t *testing.T) {
 		},
 		{
 			name: "multiple matrix responses with duplicated series, but not same samples",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"bar"},"values":[[1,"1"],[2,"2"]]}]}}`,
 				`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"bar"},"values":[[3,"3"]]}]}}`,
@@ -289,6 +323,7 @@ func TestMergeResponse(t *testing.T) {
 		},
 		{
 			name: "multiple matrix responses with duplicated series and same samples",
+			req:  defaultReq,
 			resps: []string{
 				`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"bar"},"values":[[1,"1"],[2,"2"]]}]}}`,
 				`{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"bar"},"values":[[1,"1"],[2,"2"],[3,"3"]]}]}}`,
@@ -308,7 +343,7 @@ func TestMergeResponse(t *testing.T) {
 				require.NoError(t, err)
 				resps = append(resps, dr)
 			}
-			resp, err := InstantQueryCodec.MergeResponse(context.Background(), resps...)
+			resp, err := InstantQueryCodec.MergeResponse(context.Background(), tc.req, resps...)
 			assert.Equal(t, err, tc.expectedErr)
 			if err != nil {
 				return
diff --git a/pkg/querier/tripperware/query.go b/pkg/querier/tripperware/query.go
@@ -50,7 +50,7 @@ type Codec interface {
 // Merger is used by middlewares making multiple requests to merge back all responses into a single one.
 type Merger interface {
 	// MergeResponse merges responses from multiple requests into a single Response
-	MergeResponse(context.Context, ...Response) (Response, error)
+	MergeResponse(context.Context, Request, ...Response) (Response, error)
 }
 
 // Response represents a query range response.
diff --git a/pkg/querier/tripperware/queryrange/query_range.go b/pkg/querier/tripperware/queryrange/query_range.go
@@ -127,7 +127,7 @@ func NewEmptyPrometheusResponse() *PrometheusResponse {
 	}
 }
 
-func (c prometheusCodec) MergeResponse(ctx context.Context, responses ...tripperware.Response) (tripperware.Response, error) {
+func (c prometheusCodec) MergeResponse(ctx context.Context, _ tripperware.Request, responses ...tripperware.Response) (tripperware.Response, error) {
 	sp, _ := opentracing.StartSpanFromContext(ctx, "QueryRangeResponse.MergeResponse")
 	sp.SetTag("response_count", len(responses))
 	defer sp.Finish()
diff --git a/pkg/querier/tripperware/queryrange/query_range_test.go b/pkg/querier/tripperware/queryrange/query_range_test.go
@@ -652,7 +652,7 @@ func TestMergeAPIResponses(t *testing.T) {
 			},
 		}} {
 		t.Run(tc.name, func(t *testing.T) {
-			output, err := PrometheusCodec.MergeResponse(context.Background(), tc.input...)
+			output, err := PrometheusCodec.MergeResponse(context.Background(), nil, tc.input...)
 			require.NoError(t, err)
 			require.Equal(t, tc.expected, output)
 		})
diff --git a/pkg/querier/tripperware/queryrange/results_cache.go b/pkg/querier/tripperware/queryrange/results_cache.go
@@ -407,7 +407,7 @@ func (s resultsCache) handleHit(ctx context.Context, r tripperware.Request, exte
 		return nil, nil, err
 	}
 	if len(requests) == 0 {
-		response, err := s.merger.MergeResponse(context.Background(), responses...)
+		response, err := s.merger.MergeResponse(context.Background(), r, responses...)
 		// No downstream requests so no need to write back to the cache.
 		return response, nil, err
 	}
@@ -469,7 +469,7 @@ func (s resultsCache) handleHit(ctx context.Context, r tripperware.Request, exte
 		if err != nil {
 			return nil, nil, err
 		}
-		merged, err := s.merger.MergeResponse(ctx, accumulator.Response, currentRes)
+		merged, err := s.merger.MergeResponse(ctx, r, accumulator.Response, currentRes)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -481,7 +481,7 @@ func (s resultsCache) handleHit(ctx context.Context, r tripperware.Request, exte
 		return nil, nil, err
 	}
 
-	response, err := s.merger.MergeResponse(ctx, responses...)
+	response, err := s.merger.MergeResponse(ctx, r, responses...)
 	return response, mergedExtents, err
 }
 
diff --git a/pkg/querier/tripperware/queryrange/split_by_interval.go b/pkg/querier/tripperware/queryrange/split_by_interval.go
@@ -61,7 +61,7 @@ func (s splitByInterval) Do(ctx context.Context, r tripperware.Request) (tripper
 		resps = append(resps, reqResp.Response)
 	}
 
-	response, err := s.merger.MergeResponse(ctx, resps...)
+	response, err := s.merger.MergeResponse(ctx, nil, resps...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/querier/tripperware/queryrange/split_by_interval_test.go b/pkg/querier/tripperware/queryrange/split_by_interval_test.go
@@ -266,7 +266,7 @@ func TestSplitQuery(t *testing.T) {
 }
 
 func TestSplitByDay(t *testing.T) {
-	mergedResponse, err := PrometheusCodec.MergeResponse(context.Background(), parsedResponse, parsedResponse)
+	mergedResponse, err := PrometheusCodec.MergeResponse(context.Background(), nil, parsedResponse, parsedResponse)
 	require.NoError(t, err)
 
 	mergedHTTPResponse, err := PrometheusCodec.EncodeResponse(context.Background(), mergedResponse)
diff --git a/pkg/querier/tripperware/shard_by.go b/pkg/querier/tripperware/shard_by.go
@@ -79,7 +79,7 @@ func (s shardBy) Do(ctx context.Context, r Request) (Response, error) {
 		resps = append(resps, reqResp.Response)
 	}
 
-	return s.merger.MergeResponse(ctx, resps...)
+	return s.merger.MergeResponse(ctx, r, resps...)
 }
 
 func (s shardBy) shardQuery(l log.Logger, numShards int, r Request, analysis querysharding.QueryAnalysis) []Request {

Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@ type Codec interface {`
`50`	`50`	`// Merger is used by middlewares making multiple requests to merge back all responses into a single one.`
`51`	`51`	`type Merger interface {`
`52`	`52`	`// MergeResponse merges responses from multiple requests into a single Response`
`53`		`- MergeResponse(context.Context, ...Response) (Response, error)`
	`53`	`+ MergeResponse(context.Context, Request, ...Response) (Response, error)`
`54`	`54`	`}`
`55`	`55`
`56`	`56`	`// Response represents a query range response.`
Original file line number	Diff line number	Diff line change
`@@ -127,7 +127,7 @@ func NewEmptyPrometheusResponse() *PrometheusResponse {`
`127`	`127`	`}`
`128`	`128`	`}`
`129`	`129`
`130`		`-func (c prometheusCodec) MergeResponse(ctx context.Context, responses ...tripperware.Response) (tripperware.Response, error) {`
	`130`	`+func (c prometheusCodec) MergeResponse(ctx context.Context, _ tripperware.Request, responses ...tripperware.Response) (tripperware.Response, error) {`
`131`	`131`	`sp, _ := opentracing.StartSpanFromContext(ctx, "QueryRangeResponse.MergeResponse")`
`132`	`132`	`sp.SetTag("response_count", len(responses))`
`133`	`133`	`defer sp.Finish()`
Original file line number	Diff line number	Diff line change
`@@ -407,7 +407,7 @@ func (s resultsCache) handleHit(ctx context.Context, r tripperware.Request, exte`
`407`	`407`	`return nil, nil, err`
`408`	`408`	`}`
`409`	`409`	`if len(requests) == 0 {`
`410`		`- response, err := s.merger.MergeResponse(context.Background(), responses...)`
	`410`	`+ response, err := s.merger.MergeResponse(context.Background(), r, responses...)`
`411`	`411`	`// No downstream requests so no need to write back to the cache.`
`412`	`412`	`return response, nil, err`
`413`	`413`	`}`
`@@ -469,7 +469,7 @@ func (s resultsCache) handleHit(ctx context.Context, r tripperware.Request, exte`
`469`	`469`	`if err != nil {`
`470`	`470`	`return nil, nil, err`
`471`	`471`	`}`
`472`		`- merged, err := s.merger.MergeResponse(ctx, accumulator.Response, currentRes)`
	`472`	`+ merged, err := s.merger.MergeResponse(ctx, r, accumulator.Response, currentRes)`
`473`	`473`	`if err != nil {`
`474`	`474`	`return nil, nil, err`
`475`	`475`	`}`
`@@ -481,7 +481,7 @@ func (s resultsCache) handleHit(ctx context.Context, r tripperware.Request, exte`
`481`	`481`	`return nil, nil, err`
`482`	`482`	`}`
`483`	`483`
`484`		`- response, err := s.merger.MergeResponse(ctx, responses...)`
	`484`	`+ response, err := s.merger.MergeResponse(ctx, r, responses...)`
`485`	`485`	`return response, mergedExtents, err`
`486`	`486`	`}`
`487`	`487`
Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ func (s splitByInterval) Do(ctx context.Context, r tripperware.Request) (tripper`
`61`	`61`	`resps = append(resps, reqResp.Response)`
`62`	`62`	`}`
`63`	`63`
`64`		`- response, err := s.merger.MergeResponse(ctx, resps...)`
	`64`	`+ response, err := s.merger.MergeResponse(ctx, nil, resps...)`
`65`	`65`	`if err != nil {`
`66`	`66`	`return nil, err`
`67`	`67`	`}`
Original file line number	Diff line number	Diff line change
`@@ -266,7 +266,7 @@ func TestSplitQuery(t *testing.T) {`
`266`	`266`	`}`
`267`	`267`
`268`	`268`	`func TestSplitByDay(t *testing.T) {`
`269`		`- mergedResponse, err := PrometheusCodec.MergeResponse(context.Background(), parsedResponse, parsedResponse)`
	`269`	`+ mergedResponse, err := PrometheusCodec.MergeResponse(context.Background(), nil, parsedResponse, parsedResponse)`
`270`	`270`	`require.NoError(t, err)`
`271`	`271`
`272`	`272`	`mergedHTTPResponse, err := PrometheusCodec.EncodeResponse(context.Background(), mergedResponse)`
Original file line number	Diff line number	Diff line change
`@@ -79,7 +79,7 @@ func (s shardBy) Do(ctx context.Context, r Request) (Response, error) {`
`79`	`79`	`resps = append(resps, reqResp.Response)`
`80`	`80`	`}`
`81`	`81`
`82`		`- return s.merger.MergeResponse(ctx, resps...)`
	`82`	`+ return s.merger.MergeResponse(ctx, r, resps...)`
`83`	`83`	`}`
`84`	`84`
`85`	`85`	`func (s shardBy) shardQuery(l log.Logger, numShards int, r Request, analysis querysharding.QueryAnalysis) []Request {`