support QFE instant query parse native histograms

pkg/distributor/distributor.go

@@ -627,7 +627,7 @@ func (d *Distributor) Push(ctx context.Context, req *cortexpb.WriteRequest) (*co
 	}
 	// Count the total samples, exemplars in, prior to validation or deduplication, for comparison with other metrics.
 	d.incomingSamples.WithLabelValues(userID, sampleMetricTypeFloat).Add(float64(numFloatSamples))
-	d.incomingSamples.WithLabelValues(userID, sampleMetricTypeHistogram).Add(float64(numFloatSamples))
+	d.incomingSamples.WithLabelValues(userID, sampleMetricTypeHistogram).Add(float64(numHistogramSamples))
 	d.incomingExemplars.WithLabelValues(userID).Add(float64(numExemplars))
 	// Count the total number of metadata in.
 	d.incomingMetadata.WithLabelValues(userID).Add(float64(len(req.Metadata)))

pkg/querier/tripperware/instantquery/custom.go

@@ -0,0 +1,12 @@
+package instantquery
+
+func (m *Sample) GetTimestampMs() int64 {
+	if m != nil {
+		if m.Sample != nil {
+			return m.Sample.TimestampMs
+		} else if m.Histogram != nil {
+			return m.Histogram.TimestampMs
+		}
+	}
+	return 0
+}

pkg/querier/tripperware/instantquery/instant_query.go

@@ -10,6 +10,7 @@ import (
 	"sort"
 	"strings"
 	"time"
+	"unsafe"
 
 	jsoniter "github.com/json-iterator/go"
 	"github.com/opentracing/opentracing-go"
@@ -344,7 +345,7 @@ func vectorMerge(ctx context.Context, req tripperware.Request, resps []*Promethe
 			if existingSample, ok := output[metric]; !ok {
 				output[metric] = s
 				metrics = append(metrics, metric) // Preserve the order of metric.
-			} else if existingSample.GetSample().TimestampMs < s.GetSample().TimestampMs {
+			} else if existingSample.GetTimestampMs() < s.GetTimestampMs() {
 				// Choose the latest sample if we see overlap.
 				output[metric] = s
 			}
@@ -366,11 +367,6 @@ func vectorMerge(ctx context.Context, req tripperware.Request, resps []*Promethe
 		return result, nil
 	}
 
-	type pair struct {
-		metric string
-		s      *Sample
-	}
-
 	samples := make([]*pair, 0, len(output))
 	for k, v := range output {
 		samples = append(samples, &pair{
@@ -379,13 +375,15 @@ func vectorMerge(ctx context.Context, req tripperware.Request, resps []*Promethe
 		})
 	}
 
+	// TODO: What if we have mixed float and histogram samples in the response?
+	// Then the sorting behavior is undefined. Prometheus doesn't handle it.
 	sort.Slice(samples, func(i, j int) bool {
-		// Order is determined by vector
+		// Order is determined by vector.
 		switch sortPlan {
 		case sortByValuesAsc:
-			return samples[i].s.Sample.Value < samples[j].s.Sample.Value
+			return getSortValueFromPair(samples, i) < getSortValueFromPair(samples, j)
 		case sortByValuesDesc:
-			return samples[i].s.Sample.Value > samples[j].s.Sample.Value
+			return getSortValueFromPair(samples, i) > getSortValueFromPair(samples, j)
 		}
 		return samples[i].metric < samples[j].metric
 	})
@@ -405,6 +403,22 @@ const (
 	sortByLabels     sortPlan = 3
 )
 
+type pair struct {
+	metric string
+	s      *Sample
+}
+
+// getSortValueFromPair gets the float value used for sorting from samples.
+// If float sample, use sample value. If histogram sample, use histogram sum.
+// This is the same behavior as Prometheus https://github.com/prometheus/prometheus/blob/v2.53.0/promql/functions.go#L1595.
+func getSortValueFromPair(samples []*pair, i int) float64 {
+	if samples[i].s.Histogram != nil {
+		return samples[i].s.Histogram.Histogram.Sum
+	}
+	// Impossible to have both histogram and sample nil.
+	return samples[i].s.Sample.Value
+}
+
 func sortPlanForQuery(q string) (sortPlan, error) {
 	expr, err := promqlparser.ParseExpr(q)
 	if err != nil {
@@ -534,30 +548,65 @@ func decorateWithParamName(err error, field string) error {
 	return fmt.Errorf(errTmpl, field, err)
 }
 
-// UnmarshalJSON implements json.Unmarshaler.
-func (s *Sample) UnmarshalJSON(data []byte) error {
-	var sample struct {
-		Metric labels.Labels   `json:"metric"`
-		Value  cortexpb.Sample `json:"value"`
-	}
-	if err := json.Unmarshal(data, &sample); err != nil {
-		return err
+func init() {
+	jsoniter.RegisterTypeEncoderFunc("instantquery.Sample", encodeSample, marshalJSONIsEmpty)
+	jsoniter.RegisterTypeDecoderFunc("instantquery.Sample", decodeSample)
+}
+
+func marshalJSONIsEmpty(ptr unsafe.Pointer) bool {
+	return false
+}
+
+func decodeSample(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	ss := (*Sample)(ptr)
+	for field := iter.ReadObject(); field != ""; field = iter.ReadObject() {
+		switch field {
+		case "metric":
+			metricString := iter.ReadAny().ToString()
+			lbls := labels.Labels{}
+			if err := json.UnmarshalFromString(metricString, &lbls); err != nil {
+				iter.ReportError("unmarshal Sample", err.Error())
+				return
+			}
+			ss.Labels = cortexpb.FromLabelsToLabelAdapters(lbls)
+		case "value":
+			ss.Sample = &cortexpb.Sample{}
+			cortexpb.SampleJsoniterDecode(unsafe.Pointer(ss.Sample), iter)
+		case "histogram":
+			ss.Histogram = &tripperware.SampleHistogramPair{}
+			tripperware.UnmarshalSampleHistogramPairJSON(unsafe.Pointer(ss.Histogram), iter)
+		default:
+			iter.ReportError("unmarshal Sample", fmt.Sprint("unexpected key:", field))
+			return
+		}
 	}
-	s.Labels = cortexpb.FromLabelsToLabelAdapters(sample.Metric)
-	s.Sample = sample.Value
-	return nil
 }
 
-// MarshalJSON implements json.Marshaler.
-func (s *Sample) MarshalJSON() ([]byte, error) {
-	sample := struct {
-		Metric model.Metric    `json:"metric"`
-		Value  cortexpb.Sample `json:"value"`
-	}{
-		Metric: cortexpb.FromLabelAdaptersToMetric(s.Labels),
-		Value:  s.Sample,
-	}
-	return json.Marshal(sample)
+func encodeSample(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	ss := (*Sample)(ptr)
+	stream.WriteObjectStart()
+
+	stream.WriteObjectField(`metric`)
+	lbls, err := cortexpb.FromLabelAdaptersToLabels(ss.Labels).MarshalJSON()
+	if err != nil {
+		stream.Error = err
+		return
+	}
+	stream.SetBuffer(append(stream.Buffer(), lbls...))
+
+	if ss.Sample != nil {
+		stream.WriteMore()
+		stream.WriteObjectField(`value`)
+		cortexpb.SampleJsoniterEncode(unsafe.Pointer(ss.Sample), stream)
+	}
+
+	if ss.Histogram != nil {
+		stream.WriteMore()
+		stream.WriteObjectField(`histogram`)
+		tripperware.MarshalSampleHistogramPairJSON(unsafe.Pointer(ss.Histogram), stream)
+	}
+
+	stream.WriteObjectEnd()
 }
 
 // UnmarshalJSON implements json.Unmarshaler.

pkg/querier/tripperware/instantquery/instant_query_test.go

@@ -21,6 +21,8 @@ import (
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 )
 
+const testHistogramResponse = `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"prometheus_http_request_duration_seconds","handler":"/metrics","instance":"localhost:9090","job":"prometheus"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"],[0,"0.0015060652591874421","0.001642375811042411","7"],[0,"0.001642375811042411","0.0017910235218841233","5"],[0,"0.0017910235218841233","0.001953125","13"],[0,"0.001953125","0.0021298979153618314","19"],[0,"0.0021298979153618314","0.0023226701464896895","13"],[0,"0.0023226701464896895","0.002532889755177753","13"],[0,"0.002532889755177753","0.002762135864009951","15"],[0,"0.002762135864009951","0.0030121305183748843","12"],[0,"0.0030121305183748843","0.003284751622084822","34"],[0,"0.003284751622084822","0.0035820470437682465","188"],[0,"0.0035820470437682465","0.00390625","372"],[0,"0.00390625","0.004259795830723663","400"],[0,"0.004259795830723663","0.004645340292979379","411"],[0,"0.004645340292979379","0.005065779510355506","425"],[0,"0.005065779510355506","0.005524271728019902","425"],[0,"0.005524271728019902","0.0060242610367497685","521"],[0,"0.0060242610367497685","0.006569503244169644","621"],[0,"0.006569503244169644","0.007164094087536493","593"],[0,"0.007164094087536493","0.0078125","506"],[0,"0.0078125","0.008519591661447326","458"],[0,"0.008519591661447326","0.009290680585958758","346"],[0,"0.009290680585958758","0.010131559020711013","285"],[0,"0.010131559020711013","0.011048543456039804","196"],[0,"0.011048543456039804","0.012048522073499537","129"],[0,"0.012048522073499537","0.013139006488339287","85"],[0,"0.013139006488339287","0.014328188175072986","65"],[0,"0.014328188175072986","0.015625","54"],[0,"0.015625","0.01703918332289465","53"],[0,"0.01703918332289465","0.018581361171917516","20"],[0,"0.018581361171917516","0.020263118041422026","21"],[0,"0.020263118041422026","0.022097086912079608","15"],[0,"0.022097086912079608","0.024097044146999074","11"],[0,"0.024097044146999074","0.026278012976678575","2"],[0,"0.026278012976678575","0.028656376350145972","3"],[0,"0.028656376350145972","0.03125","3"],[0,"0.04052623608284405","0.044194173824159216","2"]]}]}]}}`
+
 func TestRequest(t *testing.T) {
 	t.Parallel()
 	codec := InstantQueryCodec
@@ -185,6 +187,9 @@ func TestResponse(t *testing.T) {
 		{
 			body: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{},"value":[1,"1266464.0146205237"]}]}}`,
 		},
+		{
+			body: testHistogramResponse,
+		},
 	} {
 		tc := tc
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
@@ -260,6 +265,12 @@ func TestMergeResponse(t *testing.T) {
 			},
 			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up"},"value":[1,"1"]}]}}`,
 		},
+		{
+			name:         "duplicated histogram responses",
+			req:          defaultReq,
+			resps:        []string{testHistogramResponse, testHistogramResponse},
+			expectedResp: testHistogramResponse,
+		},
 		{
 			name: "duplicated response with stats",
 			req:  defaultReq,
@@ -278,6 +289,15 @@ func TestMergeResponse(t *testing.T) {
 			},
 			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 histogram responses",
+			req:  defaultReq,
+			resps: []string{
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"histogram":[1719528800,{"count":"1","sum":"0","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"histogram":[1719528800,{"count":"1","sum":"0","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]},{"metric":{"__name__":"up","job":"foo"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+		},
 		{
 			name: "merge two responses with sort",
 			req:  &PrometheusRequest{Query: "sort(sum by (job) (up))"},
@@ -287,6 +307,15 @@ func TestMergeResponse(t *testing.T) {
 			},
 			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 histogram responses with sort",
+			req:  &PrometheusRequest{Query: "sort(sum by (job) (up))"},
+			resps: []string{
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"histogram":[1719528880,{"count":"1","sum":"0","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"histogram":[1719528880,{"count":"1","sum":"0","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]},{"metric":{"__name__":"up","job":"foo"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+		},
 		{
 			name: "merge two responses with sort_desc",
 			req:  &PrometheusRequest{Query: "sort_desc(sum by (job) (up))"},
@@ -296,6 +325,15 @@ func TestMergeResponse(t *testing.T) {
 			},
 			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 histogram responses with sort_desc",
+			req:  &PrometheusRequest{Query: "sort_desc(sum by (job) (up))"},
+			resps: []string{
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"histogram":[1719528880,{"count":"1","sum":"0","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]},{"metric":{"__name__":"up","job":"bar"},"histogram":[1719528880,{"count":"1","sum":"0","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+		},
 		{
 			name: "merge two responses with topk",
 			req:  &PrometheusRequest{Query: "topk(10, up) by(job)"},
@@ -305,6 +343,15 @@ func TestMergeResponse(t *testing.T) {
 			},
 			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 histogram responses with topk",
+			req:  &PrometheusRequest{Query: "topk(10, up) by(job)"},
+			resps: []string{
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+				`{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"bar"},"histogram":[1719528880,{"count":"1","sum":"0","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+			},
+			expectedResp: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"up","job":"foo"},"histogram":[1719528871.898,{"count":"6342","sum":"43.31319875499995","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]},{"metric":{"__name__":"up","job":"bar"},"histogram":[1719528880,{"count":"1","sum":"0","buckets":[[0,"0.0013810679320049755","0.0015060652591874421","1"]]}]}]}}`,
+		},
 		{
 			name: "merge with warnings.",
 			req:  &PrometheusRequest{Query: "topk(10, up) by(job)"},