Ignore common label-values in ProxyResponseHeap sort function

pkg/store/proxy_heap.go

@@ -172,7 +172,7 @@ func (h *ProxyResponseHeap) Less(i, j int) bool {
 	if iResp.GetSeries() != nil && jResp.GetSeries() != nil {
 		iLbls := labelpb.ZLabelsToPromLabels(iResp.GetSeries().Labels)
 		jLbls := labelpb.ZLabelsToPromLabels(jResp.GetSeries().Labels)
-		return labels.Compare(iLbls, jLbls) < 0
+		return compareDiffLabels(iLbls, jLbls) < 0
 	} else if iResp.GetSeries() == nil && jResp.GetSeries() != nil {
 		return true
 	} else if iResp.GetSeries() != nil && jResp.GetSeries() == nil {
@@ -737,6 +737,25 @@ func newEagerRespSet(
 	return ret
 }
 
+// compareDiffLabels compares the two label sets, skipping labels with the same value in both sets.
+func compareDiffLabels(a, b labels.Labels) int {
+	a = a.Copy()
+	b = b.Copy()
+	aMap := make(map[string]string)
+	labelsToRemove := make(map[string]struct{})
+	for i := 0; i < len(a); i++ {
+		aMap[a[i].Name] = a[i].Value
+	}
+	for i := 0; i < len(b); i++ {
+		if v, ok := aMap[b[i].Name]; ok {
+			if b[i].Value == v {
+				labelsToRemove[b[i].Name] = struct{}{}
+			}
+		}
+	}
+	return labels.Compare(rmLabels(a, labelsToRemove), rmLabels(b, labelsToRemove))
+}
+
 func rmLabels(l labels.Labels, labelsToRemove map[string]struct{}) labels.Labels {
 	for i := 0; i < len(l); i++ {
 		if _, ok := labelsToRemove[l[i].Name]; !ok {

pkg/store/proxy_heap_test.go

@@ -22,6 +22,123 @@ func TestRmLabelsCornerCases(t *testing.T) {
 	}), labels.Labels{})
 }
 
+type testingRespSet struct {
+	bufferedResponse []*storepb.SeriesResponse
+	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 TestProxyResponseHeapSort(t *testing.T) {
+	for _, tcase := range []struct {
+		input []respSet
+		exp   []*storepb.SeriesResponse
+	}{
+		// Different sorted response sets
+		{
+			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", "d", "4")),
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "4", "e", "5")),
+					},
+				},
+			},
+			exp: []*storepb.SeriesResponse{
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "d", "4")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3", "d", "4")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "4", "e", "5")),
+			},
+		},
+		// Sorted response sets with different number of labels in their series
+		{
+			input: []respSet{
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+						storeSeriesResponse(t, labelsFromStrings("b", "2", "c", "3")),
+					},
+				},
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2")),
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "e", "5")),
+					},
+				},
+			},
+			exp: []*storepb.SeriesResponse{
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+				storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "e", "5")),
+				storeSeriesResponse(t, labelsFromStrings("b", "2", "c", "3")),
+			},
+		},
+		// Duplicated response sets that are sorted if you remove common label-values
+		// This is similar to response sets from stores that sort the set 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")),
+					},
+				},
+				&testingRespSet{
+					bufferedResponse: []*storepb.SeriesResponse{
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "c", "3")),
+						storeSeriesResponse(t, labelsFromStrings("a", "1", "b", "2", "c", "3")),
+					},
+				},
+			},
+			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("", 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