app/vmselect/promql: optimize binary search over big number of samples during rollup calculations

app/vmselect/promql/rollup.go

@@ -3,7 +3,6 @@ package promql
 import (
 	"fmt"
 	"math"
-	"sort"
 	"strings"
 	"sync"
 
@@ -273,10 +272,22 @@ func seekFirstTimestampIdxAfter(timestamps []int64, seekTimestamp int64, nHint i
 		return startIdx + len(timestamps)
 	}
 	// Slow path: too big len(timestamps), so use binary search.
-	i := sort.Search(len(timestamps), func(n int) bool {
+	i := binarySearchInt64(timestamps, seekTimestamp+1)
-		return n >= 0 && n < len(timestamps) && timestamps[n] > seekTimestamp
+	return startIdx + int(i)
-	})
+}
-	return startIdx + i
+
+func binarySearchInt64(a []int64, v int64) uint {
+	// Copy-pasted sort.Search from https://golang.org/src/sort/search.go?s=2246:2286#L49
+	i, j := uint(0), uint(len(a))
+	for i < j {
+		h := (i + j) >> 1
+		if h < uint(len(a)) && a[h] < v {
+			i = h + 1
+		} else {
+			j = h
+		}
+	}
+	return i
 }
 
 func getMaxPrevInterval(timestamps []int64) int64 {

app/vmselect/promql/rollup_test.go

@@ -844,6 +844,27 @@ func TestRollupFuncsNoWindow(t *testing.T) {
 	})
 }
 
+func TestRollupBigNumberOfValues(t *testing.T) {
+	const srcValuesCount = 1e4
+	rc := rollupConfig{
+		Func: rollupDefault,
+		End: srcValuesCount,
+		Step: srcValuesCount / 5,
+		Window: srcValuesCount / 4,
+	}
+	rc.Timestamps = getTimestamps(rc.Start, rc.End, rc.Step)
+	srcValues := make([]float64, srcValuesCount)
+	srcTimestamps := make([]int64, srcValuesCount)
+	for i := 0; i < srcValuesCount; i++ {
+		srcValues[i] = float64(i)
+		srcTimestamps[i] = int64(i/2)
+	}
+	values := rc.Do(nil, srcValues, srcTimestamps)
+	valuesExpected := []float64{1, 4001, 8001, 9999, nan, nan}
+	timestampsExpected := []int64{0, 2000, 4000, 6000, 8000, 10000}
+	testRowsEqual(t, values, rc.Timestamps, valuesExpected, timestampsExpected)
+}
+
 func testRowsEqual(t *testing.T, values []float64, timestamps []int64, valuesExpected []float64, timestampsExpected []int64) {
 	t.Helper()
 	if len(values) != len(valuesExpected) {

lib/storage/index_db.go

@@ -1983,7 +1983,7 @@ func (is *indexSearch) updateMetricIDsForOrSuffixWithFilter(prefix []byte, metri
 				break
 			}
 			if metricID > sf[0] {
-				n := uint64BinarySearch(sf, metricID)
+				n := binarySearchUint64(sf, metricID)
 				sf = sf[n:]
 				if len(sf) == 0 {
 					break
@@ -2006,7 +2006,7 @@ func (is *indexSearch) updateMetricIDsForOrSuffixWithFilter(prefix []byte, metri
 	return nil
 }
 
-func uint64BinarySearch(a []uint64, v uint64) uint {
+func binarySearchUint64(a []uint64, v uint64) uint {
 	// Copy-pasted sort.Search from https://golang.org/src/sort/search.go?s=2246:2286#L49
 	i, j := uint(0), uint(len(a))
 	for i < j {