perf: add searchHint to getSeriesIndex (#4286)

* perf: add searchHint to getSeriesIndex

* move searchHint to seriesIDRowsRewriter

* add comments and unit tests

* remove the second loop and needless tests

* fix unit test

Author: liaol

pkg/phlaredb/profile_store.go

@@ -516,6 +516,7 @@ type seriesIDRowsRewriter struct {
 	parquet.Rows
 	pos           int64
 	seriesIndexes rowRangesWithSeriesIndex
+	searchHint    int // speed up getSeriesIndex()
 }
 
 func (r *seriesIDRowsRewriter) SeekToRow(pos int64) error {
@@ -540,13 +541,14 @@ func (r *seriesIDRowsRewriter) ReadRows(rows []parquet.Row) (int, error) {
 	if err != nil {
 		return n, err
 	}
-
+	// sh for next call of getSeriesIndex
+	sh := r.searchHint
 	for pos, row := range rows[:n] {
 		// actual row num
 		rowNum := r.pos + int64(pos)
-		row[colIdxSeriesIndex] = parquet.ValueOf(r.seriesIndexes.getSeriesIndex(rowNum)).Level(0, 0, colIdxSeriesIndex)
+		row[colIdxSeriesIndex] = parquet.ValueOf(r.seriesIndexes.getSeriesIndex(rowNum, &sh)).Level(0, 0, colIdxSeriesIndex)
 	}
-
+	r.searchHint = sh
 	r.pos += int64(n)
 
 	return n, nil

pkg/phlaredb/profile_test.go

@@ -279,3 +279,66 @@ func TestProfileIndex_Add_OutOfOrder(t *testing.T) {
 	require.Equal(t, []int64{20, 50, 80, 100, 110}, tsOrder)
 
 }
+
+func Test_rowRangesWithSeriesIndex_getSeriesIndex(t *testing.T) {
+	testCases := []struct {
+		name              string
+		ranges            rowRangesWithSeriesIndex
+		rowNum            int64
+		searchHint        int
+		expectSeriesIndex uint32
+		expectPanic       bool
+	}{
+		{
+			name: "hit 1",
+			ranges: rowRangesWithSeriesIndex{
+				{rowRange: &rowRange{rowNum: 0, length: 5}, seriesIndex: 1},
+				{rowRange: &rowRange{rowNum: 5, length: 5}, seriesIndex: 2},
+			},
+			rowNum:            4,
+			searchHint:        0,
+			expectSeriesIndex: 1,
+		},
+		{
+			name: "hit 2",
+			ranges: rowRangesWithSeriesIndex{
+				{rowRange: &rowRange{rowNum: 0, length: 5}, seriesIndex: 1},
+				{rowRange: &rowRange{rowNum: 5, length: 5}, seriesIndex: 2},
+			},
+			rowNum:            6,
+			searchHint:        1,
+			expectSeriesIndex: 2,
+		},
+		{
+			name: "nil rowRange skipped",
+			ranges: rowRangesWithSeriesIndex{
+				{rowRange: nil, seriesIndex: 1},
+				{rowRange: &rowRange{rowNum: 10, length: 5}, seriesIndex: 2},
+			},
+			rowNum:            12,
+			searchHint:        0,
+			expectSeriesIndex: 2,
+		},
+		{
+			name:        "not found panics",
+			ranges:      rowRangesWithSeriesIndex{{rowRange: &rowRange{rowNum: 0, length: 2}, seriesIndex: 1}},
+			rowNum:      10,
+			searchHint:  0,
+			expectPanic: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			searchHint := tc.searchHint
+			if tc.expectPanic {
+				assert.Panics(t, func() {
+					_ = tc.ranges.getSeriesIndex(tc.rowNum, &searchHint)
+				})
+			} else {
+				idx := tc.ranges.getSeriesIndex(tc.rowNum, &searchHint)
+				assert.Equal(t, tc.expectSeriesIndex, idx)
+			}
+		})
+	}
+}

pkg/phlaredb/profiles.go

@@ -40,13 +40,20 @@ type rowRangeWithSeriesIndex struct {
 // those need to be strictly ordered
 type rowRangesWithSeriesIndex []rowRangeWithSeriesIndex
 
-func (s rowRangesWithSeriesIndex) getSeriesIndex(rowNum int64) uint32 {
-	for _, rg := range s {
+// getSeriesIndex returns the series index for a given row number.
+// searchHint is the hint for the index to start searching from, it should be passed to next call of this function.
+func (s rowRangesWithSeriesIndex) getSeriesIndex(rowNum int64, searchHint *int) uint32 {
+	if *searchHint < 0 || *searchHint >= len(s) {
+		*searchHint = 0
+	}
+	for i := *searchHint; i < len(s); i++ {
+		rg := s[i]
 		// it is possible that the series is not existing
 		if rg.rowRange == nil {
 			continue
 		}
 		if rg.rowNum <= rowNum && rg.rowNum+int64(rg.length) > rowNum {
+			*searchHint = i
 			return rg.seriesIndex
 		}
 	}