Ignore common label-values in ProxyResponseHeap sort function

pkg/store/proxy_heap.go

@@ -170,9 +170,13 @@ func (h *ProxyResponseHeap) Less(i, j int) bool {
 	jResp := (*h)[j].rs.At()
 
 	if iResp.GetSeries() != nil && jResp.GetSeries() != nil {
+		// Response sets are sorted before adding external labels.
+		// This comparison excludes those labels to keep the same order.
+		iStoreLbls := (*h)[i].rs.StoreLabels()
+		jStoreLbls := (*h)[j].rs.StoreLabels()
 		iLbls := labelpb.ZLabelsToPromLabels(iResp.GetSeries().Labels)
 		jLbls := labelpb.ZLabelsToPromLabels(jResp.GetSeries().Labels)
-		return labels.Compare(iLbls, jLbls) < 0
+		return labels.Compare(rmLabels(iLbls.Copy(), iStoreLbls), rmLabels(jLbls.Copy(), jStoreLbls)) < 0
 	} else if iResp.GetSeries() == nil && jResp.GetSeries() != nil {
 		return true
 	} else if iResp.GetSeries() != nil && jResp.GetSeries() == nil {
@@ -257,6 +261,10 @@ func (l *lazyRespSet) Labelset() string {
 	return labelpb.PromLabelSetsToString(l.storeLabelSets)
 }
 
+func (l *lazyRespSet) StoreLabels() map[string]struct{} {
+	return l.storeLabels
+}
+
 // lazyRespSet is a lazy storepb.SeriesSet that buffers
 // everything as fast as possible while at the same it permits
 // reading response-by-response. It blocks if there is no data
@@ -268,6 +276,7 @@ type lazyRespSet struct {
 	closeSeries    context.CancelFunc
 	storeName      string
 	storeLabelSets []labels.Labels
+	storeLabels    map[string]struct{}
 	frameTimeout   time.Duration
 	ctx            context.Context
 
@@ -374,6 +383,11 @@ func newLazyRespSet(
 		bufferedResponses:    bufferedResponses,
 		shardMatcher:         shardMatcher,
 	}
+	for _, ls := range storeLabelSets {
+		for _, l := range ls {
+			respSet.storeLabels[l.Name] = struct{}{}
+		}
+	}
 
 	go func(st string, l *lazyRespSet) {
 		bytesProcessed := 0
@@ -610,6 +624,7 @@ type eagerRespSet struct {
 
 	shardMatcher *storepb.ShardMatcher
 	removeLabels map[string]struct{}
+	storeLabels  map[string]struct{}
 
 	// Internal bookkeeping.
 	bufferedResponses []*storepb.SeriesResponse
@@ -641,6 +656,11 @@ func newEagerRespSet(
 		shardMatcher:      shardMatcher,
 		removeLabels:      removeLabels,
 	}
+	for _, ls := range st.LabelSets() {
+		for _, l := range ls {
+			ret.storeLabels[l.Name] = struct{}{}
+		}
+	}
 
 	ret.wg.Add(1)
 
@@ -811,11 +831,16 @@ func (l *eagerRespSet) Labelset() string {
 	return labelpb.PromLabelSetsToString(l.st.LabelSets())
 }
 
+func (l *eagerRespSet) StoreLabels() map[string]struct{} {
+	return l.storeLabels
+}
+
 type respSet interface {
 	Close()
 	At() *storepb.SeriesResponse
 	Next() bool
 	StoreID() string
 	Labelset() string
+	StoreLabels() map[string]struct{}
 	Empty() bool
 }

pkg/store/proxy_heap_test.go

@@ -22,6 +22,132 @@ func TestRmLabelsCornerCases(t *testing.T) {
 	}), labels.Labels{})
 }
 
+type testingRespSet struct {
+	bufferedResponse []*storepb.SeriesResponse
+	storeLabels      map[string]struct{}
+	i                int
+}
+
+func (l *testingRespSet) Close() {}
+
+func (l *testingRespSet) At() *storepb.SeriesResponse {
+	return l.bufferedResponse[l.i]
+}
+
+func (l *testingRespSet) Next() bool {
+	l.i++
+	return l.i < len(l.bufferedResponse)
+}
+
+func (l *testingRespSet) StoreID() string {
+	return ""
+}
+
+func (l *testingRespSet) Labelset() string {
+	return ""
+}
+
+func (l *testingRespSet) Empty() bool {
+	return l.i >= len(l.bufferedResponse)
+}
+
+func (l *testingRespSet) StoreLabels() map[string]struct{} {
+	return l.storeLabels
+}
+
+func TestProxyResponseHeapSort(t *testing.T) {
+	for _, tcase := range []struct {
+		title string
+		input []respSet
+		exp   []*storepb.SeriesResponse
+	}{
+		{
+			title: "merge sets with different series and common labels",
+			input: []respSet{
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3", "d", "4")),
+					},
+				},
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "4", "e", "5")),
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "d", "4")),
+					},
+				},
+			},
+			exp: []*storepb.SeriesResponse{
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3", "d", "4")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "4", "e", "5")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "d", "4")),
+			},
+		},
+		{
+			title: "merge sets with different series and labels",
+			input: []respSet{
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+						storeSeriesResponse(t, labelsFromStrings("b", "2", "c", "3")),
+						storeSeriesResponse(t, labelsFromStrings("g", "7", "h", "8", "i", "9")),
+					},
+				},
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("d", "4", "e", "5")),
+						storeSeriesResponse(t, labelsFromStrings("d", "4", "e", "5", "f", "6")),
+					},
+				},
+			},
+			exp: []*storepb.SeriesResponse{
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+				storeSeriesResponse(t, labelsFromStrings("b", "2", "c", "3")),
+				storeSeriesResponse(t, labelsFromStrings("d", "4", "e", "5")),
+				storeSeriesResponse(t, labelsFromStrings("d", "4", "e", "5", "f", "6")),
+				storeSeriesResponse(t, labelsFromStrings("g", "7", "h", "8", "i", "9")),
+			},
+		},
+		{
+			title: "merge sets that were ordered before adding external labels",
+			input: []respSet{
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+					},
+					storeLabels: map[string]struct{}{"c": {}},
+				},
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+					},
+					storeLabels: map[string]struct{}{"c": {}},
+				},
+			},
+			exp: []*storepb.SeriesResponse{
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+			},
+		},
+	} {
+		t.Run(tcase.title, func(t *testing.T) {
+			h := NewProxyResponseHeap(tcase.input...)
+			if !h.Empty() {
+				got := []*storepb.SeriesResponse{h.At()}
+				for h.Next() {
+					got = append(got, h.At())
+				}
+				testutil.Equals(t, tcase.exp, got)
+			}
+		})
+	}
+}
+
 func TestSortWithoutLabels(t *testing.T) {
 	for _, tcase := range []struct {
 		input       []*storepb.SeriesResponse