store: fix deduplication in proxy heap for series with multiple chunks

Signed-off-by: Michael Hoffmann <[email protected]>

Author: MichaHoffmann

pkg/store/proxy_heap.go

@@ -6,6 +6,7 @@ package store
 import (
 	"container/heap"
 	"context"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"sort"
@@ -111,24 +112,30 @@ func (d *dedupResponseHeap) Next() bool {
 func chainSeriesAndRemIdenticalChunks(series []*storepb.SeriesResponse) *storepb.SeriesResponse {
 	chunkDedupMap := map[uint64]*storepb.AggrChunk{}
 
+	buf := make([]byte, 8)
+	hasher := xxhash.New()
+
 	for _, s := range series {
 		for _, chk := range s.GetSeries().Chunks {
+			hasher.Reset()
 			for _, field := range []*storepb.Chunk{
 				chk.Raw, chk.Count, chk.Max, chk.Min, chk.Sum, chk.Counter,
 			} {
-				if field == nil {
-					continue
-				}
-				hash := field.Hash
-				if hash == 0 {
-					hash = xxhash.Sum64(field.Data)
-				}
-
-				if _, ok := chunkDedupMap[hash]; !ok {
-					chk := chk
-					chunkDedupMap[hash] = &chk
-					break
+				switch {
+				case field == nil:
+					binary.BigEndian.PutUint64(buf, 0)
+				case field.Hash != 0:
+					binary.BigEndian.PutUint64(buf, field.Hash)
+				default:
+					binary.BigEndian.PutUint64(buf, xxhash.Sum64(field.Data))
 				}
+				_, _ = hasher.Write(buf)
+			}
+			aggrChunkHash := hasher.Sum64()
+			if _, ok := chunkDedupMap[aggrChunkHash]; !ok {
+				chk := chk
+				chunkDedupMap[aggrChunkHash] = &chk
+				break
 			}
 		}
 	}

pkg/store/proxy_test.go

@@ -2072,6 +2072,55 @@ func TestDedupRespHeap_Deduplication(t *testing.T) {
 		testFn    func(responses []*storepb.SeriesResponse, h *dedupResponseHeap)
 		tname     string
 	}{
+		{
+			tname: "dedups identical series with multiple chunks",
+			responses: []*storepb.SeriesResponse{
+				{
+					Result: &storepb.SeriesResponse_Series{
+						Series: &storepb.Series{
+							Labels: labelpb.ZLabelsFromPromLabels(labels.FromStrings("foo", "bar")),
+							Chunks: []storepb.AggrChunk{
+								{
+									Count: &storepb.Chunk{
+										Type: storepb.Chunk_XOR,
+										Data: []byte(`count`),
+									},
+									Sum: &storepb.Chunk{
+										Type: storepb.Chunk_XOR,
+										Data: []byte(`sum`),
+									},
+								},
+							},
+						},
+					},
+				},
+				{
+					Result: &storepb.SeriesResponse_Series{
+						Series: &storepb.Series{
+							Labels: labelpb.ZLabelsFromPromLabels(labels.FromStrings("foo", "bar")),
+							Chunks: []storepb.AggrChunk{
+								{
+									Count: &storepb.Chunk{
+										Type: storepb.Chunk_XOR,
+										Data: []byte(`count`),
+									},
+									Sum: &storepb.Chunk{
+										Type: storepb.Chunk_XOR,
+										Data: []byte(`sum`),
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			testFn: func(responses []*storepb.SeriesResponse, h *dedupResponseHeap) {
+				testutil.Equals(t, true, h.Next())
+				resp := h.At()
+				testutil.Equals(t, responses[0], resp)
+				testutil.Equals(t, false, h.Next())
+			},
+		},
 		{
 			tname:     "edge case with zero responses",
 			responses: []*storepb.SeriesResponse{},