extract weight func

internal/promql/functions.go

@@ -315,49 +315,9 @@ func firstK(ev *evaluator, g seriesGroup, expr *parser.AggregateExpr) SeriesBag
 		k = len(g.bag.Data)
 	}
 	var (
-		w      = make([]float64, len(g.bag.Data))
-		nodecN int // number of non-decreasing series
+		w = g.bag.weight(ev.t.LODs)
+		x = make([]int, len(g.bag.Data))
 	)
-	for i, data := range g.bag.Data {
-		var (
-			j     int
-			acc   float64
-			prev  = -math.MaxFloat64
-			nodec = true // non-decreasing
-		)
-		for _, lod := range ev.t.LODs {
-			for m := 0; m < lod.Len; m++ {
-				v := (*data)[j+m]
-				if !math.IsNaN(v) {
-					acc += v * v * float64(lod.Step)
-					if v < prev {
-						nodec = false
-					}
-					prev = v
-				}
-			}
-			j += lod.Len
-		}
-		w[i] = acc
-		if nodec {
-			nodecN++
-		}
-	}
-	if nodecN == len(w) {
-		// all series are non-decreasing, sort by last value
-		for i, data := range g.bag.Data {
-			last := -math.MaxFloat64
-			for i := len(*data); i != 0; i-- {
-				v := (*data)[i-1]
-				if !math.IsNaN(v) {
-					last = v
-					break
-				}
-			}
-			w[i] = last
-		}
-	}
-	x := make([]int, len(g.bag.Data))
 	for i := range x {
 		x[i] = i
 	}

internal/promql/value.go

@@ -470,6 +470,53 @@ func (b *SeriesBag) trim(start, end int64) {
 	}
 }
 
+func (b *SeriesBag) weight(lods []LOD) []float64 {
+	var (
+		w      = make([]float64, len(b.Data))
+		nodecN int // number of non-decreasing series
+	)
+	for i, data := range b.Data {
+		var (
+			j     int
+			acc   float64
+			prev  = -math.MaxFloat64
+			nodec = true // non-decreasing
+		)
+		for _, lod := range lods {
+			for m := 0; m < lod.Len; m++ {
+				v := (*data)[j+m]
+				if !math.IsNaN(v) {
+					acc += v * v * float64(lod.Step)
+					if v < prev {
+						nodec = false
+					}
+					prev = v
+				}
+			}
+			j += lod.Len
+		}
+		w[i] = acc
+		if nodec {
+			nodecN++
+		}
+	}
+	if nodecN == len(w) {
+		// all series are non-decreasing, weight is a last value
+		for i, data := range b.Data {
+			last := -math.MaxFloat64
+			for i := len(*data); i != 0; i-- {
+				v := (*data)[i-1]
+				if !math.IsNaN(v) {
+					last = v
+					break
+				}
+			}
+			w[i] = last
+		}
+	}
+	return w
+}
+
 func (m *SeriesMeta) DropMetricName() {
 	m.dropTag(labels.MetricName)
 }