From 0f184affa7a270d207f6a060dc429d7a227b77bc Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Mon, 25 Nov 2019 14:01:36 +0200
Subject: [PATCH] app/vmselect/promql: optimize binary search over big number
 of samples during rollup calculations

---
 app/vmselect/promql/rollup.go      | 21 ++++++++++++++++-----
 app/vmselect/promql/rollup_test.go | 21 +++++++++++++++++++++
 lib/storage/index_db.go            |  4 ++--
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/app/vmselect/promql/rollup.go b/app/vmselect/promql/rollup.go
index cf336b2fd7..8a26bd9685 100644
--- a/app/vmselect/promql/rollup.go
+++ b/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 {
-		return n >= 0 && n < len(timestamps) && timestamps[n] > seekTimestamp
-	})
-	return startIdx + i
+	i := binarySearchInt64(timestamps, seekTimestamp+1)
+	return startIdx + int(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 {
diff --git a/app/vmselect/promql/rollup_test.go b/app/vmselect/promql/rollup_test.go
index a0acb814b7..f28593901c 100644
--- a/app/vmselect/promql/rollup_test.go
+++ b/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) {
diff --git a/lib/storage/index_db.go b/lib/storage/index_db.go
index a6e4480d34..cbab45e01b 100644
--- a/lib/storage/index_db.go
+++ b/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 {