receive: add float histogram support (thanos-io#6323)

* Added receive float histogram support Signed-off-by: Sebastian Rabenhorst <[email protected]> Fixed imports Signed-off-by: Sebastian Rabenhorst <[email protected]> * Added comments for fns copied from Prometheus Signed-off-by: Sebastian Rabenhorst <[email protected]> Improved comment Signed-off-by: Sebastian Rabenhorst <[email protected]> * Removed unnecessary if Signed-off-by: Sebastian Rabenhorst <[email protected]> * Fixed native histogram proto conversion in remote engine Signed-off-by: Sebastian Rabenhorst <[email protected]> * Removed unused histogram conversion from remote engine Signed-off-by: Sebastian Rabenhorst <[email protected]> * Fix and renaming in native_histograms_test.go Signed-off-by: Sebastian Rabenhorst <[email protected]> * Trigger Build Signed-off-by: Sebastian Rabenhorst <[email protected]> --------- Signed-off-by: Sebastian Rabenhorst <[email protected]>
vinted · Jul 27, 2023 · a48e29b · a48e29b
1 parent 089d8ce
commit a48e29b
Show file tree

Hide file tree

Showing 11 changed files with 351 additions and 236 deletions.
diff --git a/pkg/query/iter.go b/pkg/query/iter.go
@@ -178,6 +178,8 @@ func chunkEncoding(e storepb.Chunk_Encoding) chunkenc.Encoding {
 		return chunkenc.EncXOR
 	case storepb.Chunk_HISTOGRAM:
 		return chunkenc.EncHistogram
+	case storepb.Chunk_FLOAT_HISTOGRAM:
+		return chunkenc.EncFloatHistogram
 	}
 	return 255 // Invalid.
 }

diff --git a/pkg/query/remote_engine.go b/pkg/query/remote_engine.go
@@ -263,10 +263,10 @@ func (r *remoteQuery) Exec(ctx context.Context) *promql.Result {
 				F: s.Value,
 			})
 		}
-		for _, h := range ts.Histograms {
+		for _, hp := range ts.Histograms {
 			series.Histograms = append(series.Histograms, promql.HPoint{
-				T: h.Timestamp,
-				H: prompb.HistogramProtoToFloatHistogram(h),
+				T: hp.Timestamp,
+				H: prompb.FloatHistogramProtoToFloatHistogram(hp),
 			})
 		}
 		result = append(result, series)

diff --git a/pkg/receive/writer.go b/pkg/receive/writer.go
@@ -12,6 +12,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/exemplar"
+	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
@@ -168,8 +169,18 @@ func (r *Writer) Write(ctx context.Context, tenantID string, wreq *prompb.WriteR
 		}
 
 		for _, hp := range t.Histograms {
-			h := prompb.HistogramProtoToHistogram(hp)
-			ref, err = app.AppendHistogram(ref, lset, hp.Timestamp, h, nil)
+			var (
+				h  *histogram.Histogram
+				fh *histogram.FloatHistogram
+			)
+
+			if hp.IsFloatHistogram() {
+				fh = prompb.FloatHistogramProtoToFloatHistogram(hp)
+			} else {
+				h = prompb.HistogramProtoToHistogram(hp)
+			}
+
+			ref, err = app.AppendHistogram(ref, lset, hp.Timestamp, h, fh)
 			switch err {
 			case storage.ErrOutOfOrderSample:
 				numSamplesOutOfOrder++

diff --git a/pkg/receive/writer_test.go b/pkg/receive/writer_test.go
@@ -10,17 +10,17 @@ import (
 	"testing"
 	"time"
 
+	"github.com/efficientgo/core/testutil"
 	"github.com/go-kit/log"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/exemplar"
-	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
+	"github.com/prometheus/prometheus/tsdb/tsdbutil"
 
-	"github.com/efficientgo/core/testutil"
 	"github.com/thanos-io/thanos/pkg/block/metadata"
 	"github.com/thanos-io/thanos/pkg/runutil"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
@@ -251,7 +251,7 @@ func TestWriter(t *testing.T) {
 						{
 							Labels: append(lbls, labelpb.ZLabel{Name: "a", Value: "1"}, labelpb.ZLabel{Name: "b", Value: "2"}),
 							Histograms: []prompb.Histogram{
-								histogramToHistogramProto(9, testHistogram()),
+								prompb.HistogramToHistogramProto(10, tsdbutil.GenerateTestHistogram(0)),
 							},
 						},
 					},
@@ -262,7 +262,30 @@ func TestWriter(t *testing.T) {
 				{
 					Labels: append(lbls, labelpb.ZLabel{Name: "a", Value: "1"}, labelpb.ZLabel{Name: "b", Value: "2"}),
 					Histograms: []prompb.Histogram{
-						histogramToHistogramProto(10, testHistogram()),
+						prompb.HistogramToHistogramProto(10, tsdbutil.GenerateTestHistogram(0)),
+					},
+				},
+			},
+		},
+		"should succeed on float histogram with valid labels": {
+			reqs: []*prompb.WriteRequest{
+				{
+					Timeseries: []prompb.TimeSeries{
+						{
+							Labels: append(lbls, labelpb.ZLabel{Name: "a", Value: "1"}, labelpb.ZLabel{Name: "b", Value: "2"}),
+							Histograms: []prompb.Histogram{
+								prompb.FloatHistogramToHistogramProto(10, tsdbutil.GenerateTestFloatHistogram(1)),
+							},
+						},
+					},
+				},
+			},
+			expectedErr: nil,
+			expectedIngested: []prompb.TimeSeries{
+				{
+					Labels: append(lbls, labelpb.ZLabel{Name: "a", Value: "1"}, labelpb.ZLabel{Name: "b", Value: "2"}),
+					Histograms: []prompb.Histogram{
+						prompb.FloatHistogramToHistogramProto(10, tsdbutil.GenerateTestFloatHistogram(1)),
 					},
 				},
 			},
@@ -274,7 +297,7 @@ func TestWriter(t *testing.T) {
 						{
 							Labels: append(lbls, labelpb.ZLabel{Name: "a", Value: "1"}, labelpb.ZLabel{Name: "b", Value: "2"}),
 							Histograms: []prompb.Histogram{
-								histogramToHistogramProto(10, testHistogram()),
+								prompb.HistogramToHistogramProto(10, tsdbutil.GenerateTestHistogram(0)),
 							},
 						},
 					},
@@ -284,7 +307,7 @@ func TestWriter(t *testing.T) {
 						{
 							Labels: append(lbls, labelpb.ZLabel{Name: "a", Value: "1"}, labelpb.ZLabel{Name: "b", Value: "2"}),
 							Histograms: []prompb.Histogram{
-								histogramToHistogramProto(9, testHistogram()),
+								prompb.HistogramToHistogramProto(9, tsdbutil.GenerateTestHistogram(0)),
 							},
 						},
 					},
@@ -295,7 +318,7 @@ func TestWriter(t *testing.T) {
 				{
 					Labels: append(lbls, labelpb.ZLabel{Name: "a", Value: "1"}, labelpb.ZLabel{Name: "b", Value: "2"}),
 					Histograms: []prompb.Histogram{
-						histogramToHistogramProto(10, testHistogram()),
+						prompb.HistogramToHistogramProto(10, tsdbutil.GenerateTestHistogram(0)),
 					},
 				},
 			},
@@ -480,7 +503,7 @@ func generateLabelsAndSeries(numLabels int, numSeries int, generateHistograms bo
 		}
 
 		if generateHistograms {
-			ts[j].Histograms = []prompb.Histogram{histogramToHistogramProto(10, testHistogram())}
+			ts[j].Histograms = []prompb.Histogram{prompb.HistogramToHistogramProto(10, tsdbutil.GenerateTestHistogram(0))}
 			continue
 		}
 
@@ -489,42 +512,3 @@ func generateLabelsAndSeries(numLabels int, numSeries int, generateHistograms bo
 
 	return ts
 }
-
-func testHistogram() *histogram.Histogram {
-	return &histogram.Histogram{
-		Count:         5,
-		ZeroCount:     2,
-		Sum:           18.4,
-		ZeroThreshold: 0.1,
-		Schema:        1,
-		PositiveSpans: []histogram.Span{
-			{Offset: 0, Length: 2},
-			{Offset: 1, Length: 2},
-		},
-		PositiveBuckets: []int64{1, 1, -1, 0}, // counts: 1, 2, 1, 1 (total 5)
-	}
-}
-
-func histogramToHistogramProto(timestamp int64, h *histogram.Histogram) prompb.Histogram {
-	return prompb.Histogram{
-		Count:          &prompb.Histogram_CountInt{CountInt: h.Count},
-		Sum:            h.Sum,
-		Schema:         h.Schema,
-		ZeroThreshold:  h.ZeroThreshold,
-		ZeroCount:      &prompb.Histogram_ZeroCountInt{ZeroCountInt: h.ZeroCount},
-		NegativeSpans:  spansToSpansProto(h.NegativeSpans),
-		NegativeDeltas: h.NegativeBuckets,
-		PositiveSpans:  spansToSpansProto(h.PositiveSpans),
-		PositiveDeltas: h.PositiveBuckets,
-		Timestamp:      timestamp,
-	}
-}
-
-func spansToSpansProto(s []histogram.Span) []*prompb.BucketSpan {
-	spans := make([]*prompb.BucketSpan, len(s))
-	for i := 0; i < len(s); i++ {
-		spans[i] = &prompb.BucketSpan{Offset: s[i].Offset, Length: s[i].Length}
-	}
-
-	return spans
-}
diff --git a/pkg/store/storepb/prompb/custom.go b/pkg/store/storepb/prompb/custom.go
@@ -0,0 +1,9 @@
+// Copyright (c) The Thanos Authors.
+// Licensed under the Apache License 2.0.
+
+package prompb
+
+func (h Histogram) IsFloatHistogram() bool {
+	_, ok := h.GetCount().(*Histogram_CountFloat)
+	return ok
+}
diff --git a/pkg/store/storepb/prompb/samples.go b/pkg/store/storepb/prompb/samples.go
@@ -62,71 +62,124 @@ func SamplesFromPromqlSeries(series promql.Series) ([]Sample, []Histogram) {
 // HistogramProtoToHistogram extracts a (normal integer) Histogram from the
 // provided proto message. The caller has to make sure that the proto message
 // represents an interger histogram and not a float histogram.
-// Taken from https://github.com/prometheus/prometheus/blob/d33eb3ab17616a54b97d9f7791c791a79823f279/storage/remote/codec.go#L529-L542.
+// Copied from https://github.com/prometheus/prometheus/blob/0ab95536115adfe50af249d36d73674be694ca3f/storage/remote/codec.go#L626-L645
 func HistogramProtoToHistogram(hp Histogram) *histogram.Histogram {
+	if hp.IsFloatHistogram() {
+		panic("HistogramProtoToHistogram called with a float histogram")
+	}
 	return &histogram.Histogram{
-		Schema:          hp.Schema,
-		ZeroThreshold:   hp.ZeroThreshold,
-		ZeroCount:       hp.GetZeroCountInt(),
-		Count:           hp.GetCountInt(),
-		Sum:             hp.Sum,
-		PositiveSpans:   spansProtoToSpans(hp.GetPositiveSpans()),
-		PositiveBuckets: hp.GetPositiveDeltas(),
-		NegativeSpans:   spansProtoToSpans(hp.GetNegativeSpans()),
-		NegativeBuckets: hp.GetNegativeDeltas(),
+		CounterResetHint: histogram.CounterResetHint(hp.ResetHint),
+		Schema:           hp.Schema,
+		ZeroThreshold:    hp.ZeroThreshold,
+		ZeroCount:        hp.GetZeroCountInt(),
+		Count:            hp.GetCountInt(),
+		Sum:              hp.Sum,
+		PositiveSpans:    spansProtoToSpans(hp.GetPositiveSpans()),
+		PositiveBuckets:  hp.GetPositiveDeltas(),
+		NegativeSpans:    spansProtoToSpans(hp.GetNegativeSpans()),
+		NegativeBuckets:  hp.GetNegativeDeltas(),
 	}
 }
 
 // FloatHistogramToHistogramProto converts a float histogram to a protobuf type.
-// Taken from https://github.com/prometheus/prometheus/blob/d33eb3ab17616a54b97d9f7791c791a79823f279/storage/remote/codec.go#L587-L601.
-func FloatHistogramToHistogramProto(timestamp int64, fh *histogram.FloatHistogram) Histogram {
-	return Histogram{
-		Count:          &Histogram_CountFloat{CountFloat: fh.Count},
-		Sum:            fh.Sum,
-		Schema:         fh.Schema,
-		ZeroThreshold:  fh.ZeroThreshold,
-		ZeroCount:      &Histogram_ZeroCountFloat{ZeroCountFloat: fh.ZeroCount},
-		NegativeSpans:  spansToSpansProto(fh.NegativeSpans),
-		NegativeCounts: fh.NegativeBuckets,
-		PositiveSpans:  spansToSpansProto(fh.PositiveSpans),
-		PositiveCounts: fh.PositiveBuckets,
-		ResetHint:      Histogram_ResetHint(fh.CounterResetHint),
-		Timestamp:      timestamp,
+// Copied from https://github.com/prometheus/prometheus/blob/0ab95536115adfe50af249d36d73674be694ca3f/storage/remote/codec.go#L647-L667
+func FloatHistogramProtoToFloatHistogram(hp Histogram) *histogram.FloatHistogram {
+	if !hp.IsFloatHistogram() {
+		panic("FloatHistogramProtoToFloatHistogram called with an integer histogram")
+	}
+	return &histogram.FloatHistogram{
+		CounterResetHint: histogram.CounterResetHint(hp.ResetHint),
+		Schema:           hp.Schema,
+		ZeroThreshold:    hp.ZeroThreshold,
+		ZeroCount:        hp.GetZeroCountFloat(),
+		Count:            hp.GetCountFloat(),
+		Sum:              hp.Sum,
+		PositiveSpans:    spansProtoToSpans(hp.GetPositiveSpans()),
+		PositiveBuckets:  hp.GetPositiveCounts(),
+		NegativeSpans:    spansProtoToSpans(hp.GetNegativeSpans()),
+		NegativeBuckets:  hp.GetNegativeCounts(),
 	}
 }
 
 // HistogramProtoToFloatHistogram extracts a (normal integer) Histogram from the
 // provided proto message to a Float Histogram. The caller has to make sure that
 // the proto message represents an float histogram and not a integer histogram.
-// Taken from https://github.com/prometheus/prometheus/blob/d33eb3ab17616a54b97d9f7791c791a79823f279/storage/remote/codec.go#L547-L560.
+// Copied from https://github.com/prometheus/prometheus/blob/0ab95536115adfe50af249d36d73674be694ca3f/storage/remote/codec.go#L669-L688
 func HistogramProtoToFloatHistogram(hp Histogram) *histogram.FloatHistogram {
+	if hp.IsFloatHistogram() {
+		panic("HistogramProtoToFloatHistogram called with a float histogram")
+	}
 	return &histogram.FloatHistogram{
 		CounterResetHint: histogram.CounterResetHint(hp.ResetHint),
 		Schema:           hp.Schema,
 		ZeroThreshold:    hp.ZeroThreshold,
-		ZeroCount:        hp.GetZeroCountFloat(),
-		Count:            hp.GetCountFloat(),
+		ZeroCount:        float64(hp.GetZeroCountInt()),
+		Count:            float64(hp.GetCountInt()),
 		Sum:              hp.Sum,
 		PositiveSpans:    spansProtoToSpans(hp.GetPositiveSpans()),
-		PositiveBuckets:  hp.GetPositiveCounts(),
+		PositiveBuckets:  deltasToCounts(hp.GetPositiveDeltas()),
 		NegativeSpans:    spansProtoToSpans(hp.GetNegativeSpans()),
-		NegativeBuckets:  hp.GetNegativeCounts(),
+		NegativeBuckets:  deltasToCounts(hp.GetNegativeDeltas()),
 	}
 }
 
-func spansToSpansProto(s []histogram.Span) []*BucketSpan {
-	spans := make([]*BucketSpan, len(s))
+func spansProtoToSpans(s []BucketSpan) []histogram.Span {
+	spans := make([]histogram.Span, len(s))
 	for i := 0; i < len(s); i++ {
-		spans[i] = &BucketSpan{Offset: s[i].Offset, Length: s[i].Length}
+		spans[i] = histogram.Span{Offset: s[i].Offset, Length: s[i].Length}
 	}
 
 	return spans
 }
 
-func spansProtoToSpans(s []*BucketSpan) []histogram.Span {
-	spans := make([]histogram.Span, len(s))
+func deltasToCounts(deltas []int64) []float64 {
+	counts := make([]float64, len(deltas))
+	var cur float64
+	for i, d := range deltas {
+		cur += float64(d)
+		counts[i] = cur
+	}
+	return counts
+}
+
+// Copied from https://github.com/prometheus/prometheus/blob/0ab95536115adfe50af249d36d73674be694ca3f/storage/remote/codec.go#L709-L723
+func HistogramToHistogramProto(timestamp int64, h *histogram.Histogram) Histogram {
+	return Histogram{
+		Count:          &Histogram_CountInt{CountInt: h.Count},
+		Sum:            h.Sum,
+		Schema:         h.Schema,
+		ZeroThreshold:  h.ZeroThreshold,
+		ZeroCount:      &Histogram_ZeroCountInt{ZeroCountInt: h.ZeroCount},
+		NegativeSpans:  spansToSpansProto(h.NegativeSpans),
+		NegativeDeltas: h.NegativeBuckets,
+		PositiveSpans:  spansToSpansProto(h.PositiveSpans),
+		PositiveDeltas: h.PositiveBuckets,
+		ResetHint:      Histogram_ResetHint(h.CounterResetHint),
+		Timestamp:      timestamp,
+	}
+}
+
+// Copied from https://github.com/prometheus/prometheus/blob/0ab95536115adfe50af249d36d73674be694ca3f/storage/remote/codec.go#L725-L739
+func FloatHistogramToHistogramProto(timestamp int64, fh *histogram.FloatHistogram) Histogram {
+	return Histogram{
+		Count:          &Histogram_CountFloat{CountFloat: fh.Count},
+		Sum:            fh.Sum,
+		Schema:         fh.Schema,
+		ZeroThreshold:  fh.ZeroThreshold,
+		ZeroCount:      &Histogram_ZeroCountFloat{ZeroCountFloat: fh.ZeroCount},
+		NegativeSpans:  spansToSpansProto(fh.NegativeSpans),
+		NegativeCounts: fh.NegativeBuckets,
+		PositiveSpans:  spansToSpansProto(fh.PositiveSpans),
+		PositiveCounts: fh.PositiveBuckets,
+		ResetHint:      Histogram_ResetHint(fh.CounterResetHint),
+		Timestamp:      timestamp,
+	}
+}
+
+func spansToSpansProto(s []histogram.Span) []BucketSpan {
+	spans := make([]BucketSpan, len(s))
 	for i := 0; i < len(s); i++ {
-		spans[i] = histogram.Span{Offset: s[i].Offset, Length: s[i].Length}
+		spans[i] = BucketSpan{Offset: s[i].Offset, Length: s[i].Length}
 	}
 
 	return spans