writev2: Add support for nhcb.

Signed-off-by: bwplotka <bwplotka@gmail.com>

prompb/io/prometheus/write/v2/codec.go

@@ -61,7 +61,6 @@ func (h Histogram) IsFloatHistogram() bool {
 
 // ToIntHistogram returns integer Prometheus histogram from the remote implementation
 // of integer histogram. If it's a float histogram, the method returns nil.
-// TODO(bwplotka): Add support for incoming NHCB.
 func (h Histogram) ToIntHistogram() *histogram.Histogram {
 	if h.IsFloatHistogram() {
 		return nil
@@ -77,13 +76,13 @@ func (h Histogram) ToIntHistogram() *histogram.Histogram {
 		PositiveBuckets:  h.GetPositiveDeltas(),
 		NegativeSpans:    spansProtoToSpans(h.GetNegativeSpans()),
 		NegativeBuckets:  h.GetNegativeDeltas(),
+		CustomValues:     h.GetCustomValues(),
 	}
 }
 
 // ToFloatHistogram returns float Prometheus histogram from the remote implementation
 // of float histogram. If the underlying implementation is an integer histogram, a
 // conversion is performed.
-// TODO(bwplotka): Add support for incoming NHCB.
 func (h Histogram) ToFloatHistogram() *histogram.FloatHistogram {
 	if h.IsFloatHistogram() {
 		return &histogram.FloatHistogram{
@@ -97,6 +96,7 @@ func (h Histogram) ToFloatHistogram() *histogram.FloatHistogram {
 			PositiveBuckets:  h.GetPositiveCounts(),
 			NegativeSpans:    spansProtoToSpans(h.GetNegativeSpans()),
 			NegativeBuckets:  h.GetNegativeCounts(),
+			CustomValues:     h.GetCustomValues(),
 		}
 	}
 	// Conversion from integer histogram.
@@ -111,6 +111,7 @@ func (h Histogram) ToFloatHistogram() *histogram.FloatHistogram {
 		PositiveBuckets:  deltasToCounts(h.GetPositiveDeltas()),
 		NegativeSpans:    spansProtoToSpans(h.GetNegativeSpans()),
 		NegativeBuckets:  deltasToCounts(h.GetNegativeDeltas()),
+		CustomValues:     h.GetCustomValues(),
 	}
 }
 
@@ -146,6 +147,7 @@ func FromIntHistogram(timestamp int64, h *histogram.Histogram) Histogram {
 		PositiveSpans:  spansToSpansProto(h.PositiveSpans),
 		PositiveDeltas: h.PositiveBuckets,
 		ResetHint:      Histogram_ResetHint(h.CounterResetHint),
+		CustomValues:   h.CustomValues,
 		Timestamp:      timestamp,
 	}
 }
@@ -163,6 +165,7 @@ func FromFloatHistogram(timestamp int64, fh *histogram.FloatHistogram) Histogram
 		PositiveSpans:  spansToSpansProto(fh.PositiveSpans),
 		PositiveCounts: fh.PositiveBuckets,
 		ResetHint:      Histogram_ResetHint(fh.CounterResetHint),
+		CustomValues:   fh.CustomValues,
 		Timestamp:      timestamp,
 	}
 }

prompb/rwcommon/codec_test.go

@@ -113,6 +113,7 @@ func testIntHistogram() histogram.Histogram {
 			{Offset: 0, Length: 1},
 		},
 		NegativeBuckets: []int64{1, 2, -2, 1, -1, 0},
+		CustomValues:    []float64{21421, 523}, // Technically not valid nhcb, but we test if fields would be copied.
 	}
 }
 
@@ -135,22 +136,29 @@ func testFloatHistogram() histogram.FloatHistogram {
 			{Offset: 0, Length: 1},
 		},
 		NegativeBuckets: []float64{1, 3, 1, 2, 1, 1},
+		CustomValues:    []float64{21421, 523}, // Technically not valid nhcb, but we test if fields would be copied over and back.
 	}
 }
 
 func TestFromIntToFloatOrIntHistogram(t *testing.T) {
-	testIntHist := testIntHistogram()
-	testFloatHist := testFloatHistogram()
-
 	t.Run("v1", func(t *testing.T) {
-		h := prompb.FromIntHistogram(123, testIntHist.Copy())
+		// v1 does not support nhcb.
+		testIntHistWithoutNHCB := testIntHistogram()
+		testIntHistWithoutNHCB.CustomValues = nil
+		testFloatHistWithoutNHCB := testFloatHistogram()
+		testFloatHistWithoutNHCB.CustomValues = nil
+
+		h := prompb.FromIntHistogram(123, &testIntHistWithoutNHCB)
 		require.False(t, h.IsFloatHistogram())
 		require.Equal(t, int64(123), h.Timestamp)
-		require.Equal(t, testIntHist, *h.ToIntHistogram())
+		require.Equal(t, testIntHistWithoutNHCB, *h.ToIntHistogram())
-		require.Equal(t, testFloatHist, *h.ToFloatHistogram())
+		require.Equal(t, testFloatHistWithoutNHCB, *h.ToFloatHistogram())
 	})
 	t.Run("v2", func(t *testing.T) {
-		h := writev2.FromIntHistogram(123, testIntHist.Copy())
+		testIntHist := testIntHistogram()
+		testFloatHist := testFloatHistogram()
+
+		h := writev2.FromIntHistogram(123, &testIntHist)
 		require.False(t, h.IsFloatHistogram())
 		require.Equal(t, int64(123), h.Timestamp)
 		require.Equal(t, testIntHist, *h.ToIntHistogram())
@@ -159,17 +167,21 @@ func TestFromIntToFloatOrIntHistogram(t *testing.T) {
 }
 
 func TestFromFloatToFloatHistogram(t *testing.T) {
-	testFloatHist := testFloatHistogram()
-
 	t.Run("v1", func(t *testing.T) {
-		h := prompb.FromFloatHistogram(123, testFloatHist.Copy())
+		// v1 does not support nhcb.
+		testFloatHistWithoutNHCB := testFloatHistogram()
+		testFloatHistWithoutNHCB.CustomValues = nil
+
+		h := prompb.FromFloatHistogram(123, &testFloatHistWithoutNHCB)
 		require.True(t, h.IsFloatHistogram())
 		require.Equal(t, int64(123), h.Timestamp)
 		require.Nil(t, h.ToIntHistogram())
-		require.Equal(t, testFloatHist, *h.ToFloatHistogram())
+		require.Equal(t, testFloatHistWithoutNHCB, *h.ToFloatHistogram())
 	})
 	t.Run("v2", func(t *testing.T) {
-		h := writev2.FromFloatHistogram(123, testFloatHist.Copy())
+		testFloatHist := testFloatHistogram()
+
+		h := writev2.FromFloatHistogram(123, &testFloatHist)
 		require.True(t, h.IsFloatHistogram())
 		require.Equal(t, int64(123), h.Timestamp)
 		require.Nil(t, h.ToIntHistogram())