Optimize PromQL aggregations (#4248)

* Compute hash of label subsets without creating a LabelSet first.

Signed-off-by: Alin Sinpalean <alin.sinpalean@gmail.com>

pkg/labels/labels.go

@@ -94,6 +94,47 @@ func (ls Labels) Hash() uint64 {
 	return xxhash.Sum64(b)
 }
 
+// HashForLabels returns a hash value for the labels matching the provided names.
+func (ls Labels) HashForLabels(names ...string) uint64 {
+	b := make([]byte, 0, 1024)
+
+	for _, v := range ls {
+		for _, n := range names {
+			if v.Name == n {
+				b = append(b, v.Name...)
+				b = append(b, sep)
+				b = append(b, v.Value...)
+				b = append(b, sep)
+				break
+			}
+		}
+	}
+	return xxhash.Sum64(b)
+}
+
+// HashWithoutLabels returns a hash value for all labels except those matching
+// the provided names.
+func (ls Labels) HashWithoutLabels(names ...string) uint64 {
+	b := make([]byte, 0, 1024)
+
+Outer:
+	for _, v := range ls {
+		if v.Name == MetricName {
+			continue
+		}
+		for _, n := range names {
+			if v.Name == n {
+				continue Outer
+			}
+		}
+		b = append(b, v.Name...)
+		b = append(b, sep)
+		b = append(b, v.Value...)
+		b = append(b, sep)
+	}
+	return xxhash.Sum64(b)
+}
+
 // Copy returns a copy of the labels.
 func (ls Labels) Copy() Labels {
 	res := make(Labels, len(ls))

promql/bench_test.go

@@ -131,6 +131,22 @@ func BenchmarkRangeQuery(b *testing.B) {
 		{
 			expr: "label_join(a_X, 'l2', '-', 'l', 'l')",
 		},
+		// Simple aggregations.
+		{
+			expr: "sum(a_X)",
+		},
+		{
+			expr: "sum without (l)(h_X)",
+		},
+		{
+			expr: "sum without (le)(h_X)",
+		},
+		{
+			expr: "sum by (l)(h_X)",
+		},
+		{
+			expr: "sum by (le)(h_X)",
+		},
 		// Combinations.
 		{
 			expr: "rate(a_X[1m]) + rate(b_X[1m])",

promql/engine.go

@@ -1261,49 +1261,16 @@ func (ev *evaluator) VectorBinop(op ItemType, lhs, rhs Vector, matching *VectorM
 	return enh.out
 }
 
-func hashWithoutLabels(lset labels.Labels, names ...string) uint64 {
-	cm := make(labels.Labels, 0, len(lset))
-
-Outer:
-	for _, l := range lset {
-		for _, n := range names {
-			if n == l.Name {
-				continue Outer
-			}
-		}
-		if l.Name == labels.MetricName {
-			continue
-		}
-		cm = append(cm, l)
-	}
-
-	return cm.Hash()
-}
-
-func hashForLabels(lset labels.Labels, names ...string) uint64 {
-	cm := make(labels.Labels, 0, len(names))
-
-	for _, l := range lset {
-		for _, n := range names {
-			if l.Name == n {
-				cm = append(cm, l)
-				break
-			}
-		}
-	}
-	return cm.Hash()
-}
-
 // signatureFunc returns a function that calculates the signature for a metric
 // ignoring the provided labels. If on, then the given labels are only used instead.
 func signatureFunc(on bool, names ...string) func(labels.Labels) uint64 {
 	// TODO(fabxc): ensure names are sorted and then use that and sortedness
 	// of labels by names to speed up the operations below.
 	// Alternatively, inline the hashing and don't build new label sets.
 	if on {
-		return func(lset labels.Labels) uint64 { return hashForLabels(lset, names...) }
+		return func(lset labels.Labels) uint64 { return lset.HashForLabels(names...) }
 	}
-	return func(lset labels.Labels) uint64 { return hashWithoutLabels(lset, names...) }
+	return func(lset labels.Labels) uint64 { return lset.HashWithoutLabels(names...) }
 }
 
 // resultMetric returns the metric for the given sample(s) based on the Vector
@@ -1504,24 +1471,21 @@ func (ev *evaluator) aggregation(op ItemType, grouping []string, without bool, p
 	}
 
 	for _, s := range vec {
-		lb := labels.NewBuilder(s.Metric)
+		metric := s.Metric
 
-		if without {
-			lb.Del(grouping...)
-			lb.Del(labels.MetricName)
-		}
 		if op == itemCountValues {
+			lb := labels.NewBuilder(metric)
 			lb.Set(valueLabel, strconv.FormatFloat(s.V, 'f', -1, 64))
+			metric = lb.Labels()
 		}
 
 		var (
 			groupingKey uint64
-			metric      = lb.Labels()
 		)
 		if without {
-			groupingKey = metric.Hash()
+			groupingKey = metric.HashWithoutLabels(grouping...)
 		} else {
-			groupingKey = hashForLabels(metric, grouping...)
+			groupingKey = metric.HashForLabels(grouping...)
 		}
 
 		group, ok := result[groupingKey]
@@ -1530,13 +1494,16 @@ func (ev *evaluator) aggregation(op ItemType, grouping []string, without bool, p
 			var m labels.Labels
 
 			if without {
-				m = metric
+				lb := labels.NewBuilder(metric)
+				lb.Del(grouping...)
+				lb.Del(labels.MetricName)
+				m = lb.Labels()
 			} else {
 				m = make(labels.Labels, 0, len(grouping))
 				for _, l := range metric {
 					for _, n := range grouping {
 						if l.Name == n {
-							m = append(m, labels.Label{Name: n, Value: l.Value})
+							m = append(m, l)
 							break
 						}
 					}