Merge pull request #15482 from charleskorn/charleskorn/nh-shared-customvalues

tsdb/chunkenc: don't share custom value slices between histograms

tsdb/chunkenc/float_histogram.go

@@ -976,6 +976,7 @@ func (it *floatHistogramIterator) Reset(b []byte) {
 		it.atFloatHistogramCalled = false
 		it.pBuckets, it.nBuckets = nil, nil
 		it.pSpans, it.nSpans = nil, nil
+		it.customValues = nil
 	} else {
 		it.pBuckets, it.nBuckets = it.pBuckets[:0], it.nBuckets[:0]
 	}
@@ -1071,7 +1072,7 @@ func (it *floatHistogramIterator) Next() ValueType {
 	// The case for the 2nd sample with single deltas is implicitly handled correctly with the double delta code,
 	// so we don't need a separate single delta logic for the 2nd sample.
 
-	// Recycle bucket and span slices that have not been returned yet. Otherwise, copy them.
+	// Recycle bucket, span and custom value slices that have not been returned yet. Otherwise, copy them.
 	// We can always recycle the slices for leading and trailing bits as they are
 	// never returned to the caller.
 	if it.atFloatHistogramCalled {
@@ -1104,6 +1105,13 @@ func (it *floatHistogramIterator) Next() ValueType {
 		} else {
 			it.nSpans = nil
 		}
+		if len(it.customValues) > 0 {
+			newCustomValues := make([]float64, len(it.customValues))
+			copy(newCustomValues, it.customValues)
+			it.customValues = newCustomValues
+		} else {
+			it.customValues = nil
+		}
 	}
 
 	tDod, err := readVarbitInt(&it.br)

tsdb/chunkenc/float_histogram_test.go

@@ -1368,6 +1368,52 @@ func TestFloatHistogramUniqueSpansAfterNext(t *testing.T) {
 	require.NotSame(t, &rh1.NegativeSpans[0], &rh2.NegativeSpans[0], "NegativeSpans should be unique between histograms")
 }
 
+func TestFloatHistogramUniqueCustomValuesAfterNext(t *testing.T) {
+	// Create two histograms with the same schema and custom values.
+	h1 := &histogram.FloatHistogram{
+		Schema: -53,
+		Count:  10,
+		Sum:    15.0,
+		PositiveSpans: []histogram.Span{
+			{Offset: 0, Length: 2},
+			{Offset: 1, Length: 2},
+		},
+		PositiveBuckets: []float64{1, 2, 3, 4},
+		CustomValues:    []float64{10, 11, 12, 13},
+	}
+
+	h2 := h1.Copy()
+
+	// Create a chunk and append both histograms.
+	c := NewFloatHistogramChunk()
+	app, err := c.Appender()
+	require.NoError(t, err)
+
+	_, _, _, err = app.AppendFloatHistogram(nil, 0, h1, false)
+	require.NoError(t, err)
+
+	_, _, _, err = app.AppendFloatHistogram(nil, 1, h2, false)
+	require.NoError(t, err)
+
+	// Create an iterator and advance to the first histogram.
+	it := c.Iterator(nil)
+	require.Equal(t, ValFloatHistogram, it.Next())
+	_, rh1 := it.AtFloatHistogram(nil)
+
+	// Advance to the second histogram and retrieve it.
+	require.Equal(t, ValFloatHistogram, it.Next())
+	_, rh2 := it.AtFloatHistogram(nil)
+
+	require.Equal(t, rh1.PositiveSpans, h1.PositiveSpans, "Returned positive spans are as expected")
+	require.Equal(t, rh1.CustomValues, h1.CustomValues, "Returned custom values are as expected")
+	require.Equal(t, rh2.PositiveSpans, h1.PositiveSpans, "Returned positive spans are as expected")
+	require.Equal(t, rh2.CustomValues, h1.CustomValues, "Returned custom values are as expected")
+
+	// Check that the spans and custom values for h1 and h2 are unique slices.
+	require.NotSame(t, &rh1.PositiveSpans[0], &rh2.PositiveSpans[0], "PositiveSpans should be unique between histograms")
+	require.NotSame(t, &rh1.CustomValues[0], &rh2.CustomValues[0], "CustomValues should be unique between histograms")
+}
+
 func assertFirstFloatHistogramSampleHint(t *testing.T, chunk Chunk, expected histogram.CounterResetHint) {
 	it := chunk.Iterator(nil)
 	require.Equal(t, ValFloatHistogram, it.Next())

tsdb/chunkenc/histogram.go

@@ -1075,6 +1075,7 @@ func (it *histogramIterator) Reset(b []byte) {
 		it.atHistogramCalled = false
 		it.pBuckets, it.nBuckets = nil, nil
 		it.pSpans, it.nSpans = nil, nil
+		it.customValues = nil
 	} else {
 		it.pBuckets = it.pBuckets[:0]
 		it.nBuckets = it.nBuckets[:0]
@@ -1082,6 +1083,7 @@ func (it *histogramIterator) Reset(b []byte) {
 	if it.atFloatHistogramCalled {
 		it.atFloatHistogramCalled = false
 		it.pFloatBuckets, it.nFloatBuckets = nil, nil
+		it.customValues = nil
 	} else {
 		it.pFloatBuckets = it.pFloatBuckets[:0]
 		it.nFloatBuckets = it.nFloatBuckets[:0]
@@ -1187,8 +1189,7 @@ func (it *histogramIterator) Next() ValueType {
 	// The case for the 2nd sample with single deltas is implicitly handled correctly with the double delta code,
 	// so we don't need a separate single delta logic for the 2nd sample.
 
-	// Recycle bucket and span slices that have not been returned yet. Otherwise, copy them.
-	// copy them.
+	// Recycle bucket, span and custom value slices that have not been returned yet. Otherwise, copy them.
 	if it.atFloatHistogramCalled || it.atHistogramCalled {
 		if len(it.pSpans) > 0 {
 			newSpans := make([]histogram.Span, len(it.pSpans))
@@ -1204,6 +1205,13 @@ func (it *histogramIterator) Next() ValueType {
 		} else {
 			it.nSpans = nil
 		}
+		if len(it.customValues) > 0 {
+			newCustomValues := make([]float64, len(it.customValues))
+			copy(newCustomValues, it.customValues)
+			it.customValues = newCustomValues
+		} else {
+			it.customValues = nil
+		}
 	}
 
 	if it.atHistogramCalled {

tsdb/chunkenc/histogram_test.go

@@ -1497,7 +1497,7 @@ func TestHistogramAppendOnlyErrors(t *testing.T) {
 	})
 }
 
-func TestHistogramUniqueSpansAfterNext(t *testing.T) {
+func TestHistogramUniqueSpansAfterNextWithAtHistogram(t *testing.T) {
 	// Create two histograms with the same schema and spans.
 	h1 := &histogram.Histogram{
 		Schema:        1,
@@ -1599,6 +1599,98 @@ func TestHistogramUniqueSpansAfterNextWithAtFloatHistogram(t *testing.T) {
 	require.NotSame(t, &rh1.NegativeSpans[0], &rh2.NegativeSpans[0], "NegativeSpans should be unique between histograms")
 }
 
+func TestHistogramUniqueCustomValuesAfterNextWithAtHistogram(t *testing.T) {
+	// Create two histograms with the same schema and custom values.
+	h1 := &histogram.Histogram{
+		Schema: -53,
+		Count:  10,
+		Sum:    15.0,
+		PositiveSpans: []histogram.Span{
+			{Offset: 0, Length: 2},
+			{Offset: 1, Length: 2},
+		},
+		PositiveBuckets: []int64{1, 2, 3, 4},
+		CustomValues:    []float64{10, 11, 12, 13},
+	}
+
+	h2 := h1.Copy()
+
+	// Create a chunk and append both histograms.
+	c := NewHistogramChunk()
+	app, err := c.Appender()
+	require.NoError(t, err)
+
+	_, _, _, err = app.AppendHistogram(nil, 0, h1, false)
+	require.NoError(t, err)
+
+	_, _, _, err = app.AppendHistogram(nil, 1, h2, false)
+	require.NoError(t, err)
+
+	// Create an iterator and advance to the first histogram.
+	it := c.Iterator(nil)
+	require.Equal(t, ValHistogram, it.Next())
+	_, rh1 := it.AtHistogram(nil)
+
+	// Advance to the second histogram and retrieve it.
+	require.Equal(t, ValHistogram, it.Next())
+	_, rh2 := it.AtHistogram(nil)
+
+	require.Equal(t, rh1.PositiveSpans, h1.PositiveSpans, "Returned positive spans are as expected")
+	require.Equal(t, rh1.CustomValues, h1.CustomValues, "Returned custom values are as expected")
+	require.Equal(t, rh2.PositiveSpans, h1.PositiveSpans, "Returned positive spans are as expected")
+	require.Equal(t, rh2.CustomValues, h1.CustomValues, "Returned custom values are as expected")
+
+	// Check that the spans and custom values for h1 and h2 are unique slices.
+	require.NotSame(t, &rh1.PositiveSpans[0], &rh2.PositiveSpans[0], "PositiveSpans should be unique between histograms")
+	require.NotSame(t, &rh1.CustomValues[0], &rh2.CustomValues[0], "CustomValues should be unique between histograms")
+}
+
+func TestHistogramUniqueCustomValuesAfterNextWithAtFloatHistogram(t *testing.T) {
+	// Create two histograms with the same schema and custom values.
+	h1 := &histogram.Histogram{
+		Schema: -53,
+		Count:  10,
+		Sum:    15.0,
+		PositiveSpans: []histogram.Span{
+			{Offset: 0, Length: 2},
+			{Offset: 1, Length: 2},
+		},
+		PositiveBuckets: []int64{1, 2, 3, 4},
+		CustomValues:    []float64{10, 11, 12, 13},
+	}
+
+	h2 := h1.Copy()
+
+	// Create a chunk and append both histograms.
+	c := NewHistogramChunk()
+	app, err := c.Appender()
+	require.NoError(t, err)
+
+	_, _, _, err = app.AppendHistogram(nil, 0, h1, false)
+	require.NoError(t, err)
+
+	_, _, _, err = app.AppendHistogram(nil, 1, h2, false)
+	require.NoError(t, err)
+
+	// Create an iterator and advance to the first histogram.
+	it := c.Iterator(nil)
+	require.Equal(t, ValHistogram, it.Next())
+	_, rh1 := it.AtFloatHistogram(nil)
+
+	// Advance to the second histogram and retrieve it.
+	require.Equal(t, ValHistogram, it.Next())
+	_, rh2 := it.AtFloatHistogram(nil)
+
+	require.Equal(t, rh1.PositiveSpans, h1.PositiveSpans, "Returned positive spans are as expected")
+	require.Equal(t, rh1.CustomValues, h1.CustomValues, "Returned custom values are as expected")
+	require.Equal(t, rh2.PositiveSpans, h1.PositiveSpans, "Returned positive spans are as expected")
+	require.Equal(t, rh2.CustomValues, h1.CustomValues, "Returned custom values are as expected")
+
+	// Check that the spans and custom values for h1 and h2 are unique slices.
+	require.NotSame(t, &rh1.PositiveSpans[0], &rh2.PositiveSpans[0], "PositiveSpans should be unique between histograms")
+	require.NotSame(t, &rh1.CustomValues[0], &rh2.CustomValues[0], "CustomValues should be unique between histograms")
+}
+
 func BenchmarkAppendable(b *testing.B) {
 	// Create a histogram with a bunch of spans and buckets.
 	const (