Support sharding subqueries in instant query QF

CHANGELOG.md

@@ -64,6 +64,7 @@
 * [FEATURE] Distributor: Added a new limit `-validation.max-labels-size-bytes` allowing to limit the combined size of labels for each timeseries. #4848
 * [FEATURE] Storage/Bucket: Added `-*.s3.bucket-lookup-type` allowing to configure the s3 bucket lookup type. #4794
 * [FEATURE] QueryFrontend: Implement experimental vertical sharding at query frontend for range/instant queries. #4863
+* [FEATURE] QueryFrontend: Support vertical sharding for subqueries. #4955
 * [FEATURE] Querier: Added a new limit `-querier.max-fetched-data-bytes-per-query` allowing to limit the maximum size of all data in bytes that a query can fetch from each ingester and storage. #4854
 * [FEATURE] Added 2 flags `-alertmanager.alertmanager-client.grpc-compression` and `-querier.store-gateway-client.grpc-compression` to configure compression methods for grpc clients. #4889
 * [BUGFIX] Storage/Bucket: Enable AWS SDK for go authentication for s3 to fix IMDSv1 authentication. #4897

go.mod

@@ -48,7 +48,7 @@ require (
 	github.com/spf13/afero v1.6.0
 	github.com/stretchr/testify v1.8.0
 	github.com/thanos-io/objstore v0.0.0-20221006135717-79dcec7fe604
-	github.com/thanos-io/thanos v0.27.0-rc.0.0.20221020152244-d5744e794d97
+	github.com/thanos-io/thanos v0.27.0-rc.0.0.20221102162214-50b41561b178
 	github.com/uber/jaeger-client-go v2.30.0+incompatible
 	github.com/weaveworks/common v0.0.0-20220706100410-67d27ed40fae
 	go.etcd.io/etcd/api/v3 v3.5.4

go.sum

@@ -927,8 +927,8 @@ github.com/thanos-community/galaxycache v0.0.0-20211122094458-3a32041a1f1e h1:f1
 github.com/thanos-community/galaxycache v0.0.0-20211122094458-3a32041a1f1e/go.mod h1:jXcofnrSln/cLI6/dhlBxPQZEEQHVPCcFaH75M+nSzM=
 github.com/thanos-io/objstore v0.0.0-20221006135717-79dcec7fe604 h1:9dceDSKKsLWNHjrMpyzK1t7eVcAZv9Dp3FX+uokUS2Y=
 github.com/thanos-io/objstore v0.0.0-20221006135717-79dcec7fe604/go.mod h1:Vx5dZs9ElxEhNLnum/OgB0pNTqNdI2zdXL82BeJr3T4=
-github.com/thanos-io/thanos v0.27.0-rc.0.0.20221020152244-d5744e794d97 h1:FPFrA/tx2XGRGuIi3FArTosESoZc8fJZDOCug2hhYvc=
-github.com/thanos-io/thanos v0.27.0-rc.0.0.20221020152244-d5744e794d97/go.mod h1:hjmmdeaAWtZFrIBFR3tbzhPEajIrg+A3lG55bYAXZYA=
+github.com/thanos-io/thanos v0.27.0-rc.0.0.20221102162214-50b41561b178 h1:rZTdKdZyyG+6Yi50mOkRmw7BFmZhdnFr/N+8o+Enadw=
+github.com/thanos-io/thanos v0.27.0-rc.0.0.20221102162214-50b41561b178/go.mod h1:odqdxSO+o/UaVgNpdkYYaQUW/JpT7LByXyZmxoe6uoc=
 github.com/themihai/gomemcache v0.0.0-20180902122335-24332e2d58ab h1:7ZR3hmisBWw77ZpO1/o86g+JV3VKlk3d48jopJxzTjU=
 github.com/themihai/gomemcache v0.0.0-20180902122335-24332e2d58ab/go.mod h1:eheTFp954zcWZXCU8d0AT76ftsQOTo4DTqkN/h3k1MY=
 github.com/tidwall/pretty v0.0.0-20180105212114-65a9db5fad51/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=

pkg/querier/tripperware/instantquery/instant_query.go

@@ -253,26 +253,38 @@ func (instantQueryCodec) MergeResponse(responses ...tripperware.Response) (tripp
 		promResponses = append(promResponses, resp.(*PrometheusInstantQueryResponse))
 	}
 
-	var r isPrometheusInstantQueryResult_Result
+	var data PrometheusInstantQueryData
 	// For now, we only shard queries that returns a vector.
 	switch promResponses[0].Data.ResultType {
 	case model.ValVector.String():
-		r = &PrometheusInstantQueryResult_Vector{
-			Vector: vectorMerge(promResponses),
+		data = PrometheusInstantQueryData{
+			ResultType: model.ValVector.String(),
+			Result: PrometheusInstantQueryResult{
+				Result: &PrometheusInstantQueryResult_Vector{
+					Vector: vectorMerge(promResponses),
+				},
+			},
+			Stats: statsMerge(promResponses),
+		}
+	case model.ValMatrix.String():
+		data = PrometheusInstantQueryData{
+			ResultType: model.ValMatrix.String(),
+			Result: PrometheusInstantQueryResult{
+				Result: &PrometheusInstantQueryResult_Matrix{
+					Matrix: &Matrix{
+						SampleStreams: matrixMerge(promResponses),
+					},
+				},
+			},
+			Stats: statsMerge(promResponses),
 		}
 	default:
 		return nil, fmt.Errorf("unexpected result type on instant query: %s", promResponses[0].Data.ResultType)
 	}
 
 	res := &PrometheusInstantQueryResponse{
 		Status: queryrange.StatusSuccess,
-		Data: PrometheusInstantQueryData{
-			ResultType: model.ValVector.String(),
-			Result: PrometheusInstantQueryResult{
-				Result: r,
-			},
-			Stats: statsMerge(promResponses),
-		},
+		Data:   data,
 	}
 	return res, nil
 }
@@ -324,6 +336,32 @@ func vectorMerge(resps []*PrometheusInstantQueryResponse) *Vector {
 	return result
 }
 
+func matrixMerge(resps []*PrometheusInstantQueryResponse) []tripperware.SampleStream {
+	output := make(map[string]tripperware.SampleStream)
+	for _, resp := range resps {
+		if resp == nil {
+			continue
+		}
+		if resp.Data.Result.GetMatrix() == nil {
+			continue
+		}
+		tripperware.MergeSampleStreams(output, resp.Data.Result.GetMatrix().GetSampleStreams())
+	}
+
+	keys := make([]string, 0, len(output))
+	for key := range output {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+
+	result := make([]tripperware.SampleStream, 0, len(output))
+	for _, key := range keys {
+		result = append(result, output[key])
+	}
+
+	return result
+}
+
 // NewEmptyPrometheusInstantQueryResponse returns an empty successful Prometheus query range response.
 func NewEmptyPrometheusInstantQueryResponse() *PrometheusInstantQueryResponse {
 	return &PrometheusInstantQueryResponse{
@@ -425,6 +463,18 @@ func (s *PrometheusInstantQueryData) UnmarshalJSON(data []byte) error {
 				Samples: result.Samples,
 			}},
 		}
+	case model.ValMatrix.String():
+		var result struct {
+			SampleStreams []tripperware.SampleStream `json:"result"`
+		}
+		if err := json.Unmarshal(data, &result); err != nil {
+			return err
+		}
+		s.Result = PrometheusInstantQueryResult{
+			Result: &PrometheusInstantQueryResult_Matrix{Matrix: &Matrix{
+				SampleStreams: result.SampleStreams,
+			}},
+		}
 	default:
 		s.Result = PrometheusInstantQueryResult{
 			Result: &PrometheusInstantQueryResult_RawBytes{data},
@@ -447,6 +497,17 @@ func (s *PrometheusInstantQueryData) MarshalJSON() ([]byte, error) {
 			Stats:      s.Stats,
 		}
 		return json.Marshal(res)
+	case model.ValMatrix.String():
+		res := struct {
+			ResultType string                               `json:"resultType"`
+			Data       []tripperware.SampleStream           `json:"result"`
+			Stats      *tripperware.PrometheusResponseStats `json:"stats,omitempty"`
+		}{
+			ResultType: s.ResultType,
+			Data:       s.Result.GetMatrix().SampleStreams,
+			Stats:      s.Stats,
+		}
+		return json.Marshal(res)
 	default:
 		return s.Result.GetRawBytes(), nil
 	}

pkg/querier/tripperware/instantquery/instant_query_test.go

@@ -206,6 +206,37 @@ func TestMergeResponse(t *testing.T) {
 			},
 			expectedErr: fmt.Errorf("unexpected result type on instant query: %s", "string"),
 		},
+		{
+			name: "single matrix response",
+			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",
+			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"]]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"bar"},"values":[[1,"1"],[2,"2"]]},{"metric":{"__name__":"foo"},"values":[[3,"3"],[4,"4"]]}]}}`,
+		},
+		{
+			name: "multiple matrix responses with duplicated series, but not same samples",
+			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"]]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"bar"},"values":[[1,"1"],[2,"2"],[3,"3"]]}]}}`,
+		},
+		{
+			name: "multiple matrix responses with duplicated series and same samples",
+			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"]]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"__name__":"bar"},"values":[[1,"1"],[2,"2"],[3,"3"]]}]}}`,
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			var resps []tripperware.Response