Binary search based table picker (#983)

Use binary search to pick tables instead of simple looping.
This seems to improve the performance of key iterator 
creation close to seek.

Signed-off-by: பாலாஜி <[email protected]>

Author: balaji

iterator.go

@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"fmt"
 	"hash/crc32"
+	"sort"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -339,20 +340,23 @@ type IteratorOptions struct {
 	InternalAccess bool // Used to allow internal access to badger keys.
 }
 
+func (opt *IteratorOptions) compareToPrefix(key []byte) int {
+	// We should compare key without timestamp. For example key - a[TS] might be > "aa" prefix.
+	key = y.ParseKey(key)
+	if len(key) > len(opt.Prefix) {
+		key = key[:len(opt.Prefix)]
+	}
+	return bytes.Compare(key, opt.Prefix)
+}
+
 func (opt *IteratorOptions) pickTable(t table.TableInterface) bool {
 	if len(opt.Prefix) == 0 {
 		return true
 	}
-	trim := func(key []byte) []byte {
-		if len(key) > len(opt.Prefix) {
-			return key[:len(opt.Prefix)]
-		}
-		return key
-	}
-	if bytes.Compare(trim(y.ParseKey(t.Smallest())), opt.Prefix) > 0 {
+	if opt.compareToPrefix(t.Smallest()) > 0 {
 		return false
 	}
-	if bytes.Compare(trim(y.ParseKey(t.Biggest())), opt.Prefix) < 0 {
+	if opt.compareToPrefix(t.Biggest()) < 0 {
 		return false
 	}
 	// Bloom filter lookup would only work if opt.Prefix does NOT have the read
@@ -363,6 +367,49 @@ func (opt *IteratorOptions) pickTable(t table.TableInterface) bool {
 	return true
 }
 
+// pickTables picks the necessary table for the iterator. This function also assumes
+// that the tables are sorted in the right order.
+func (opt *IteratorOptions) pickTables(all []*table.Table) []*table.Table {
+	if len(opt.Prefix) == 0 {
+		out := make([]*table.Table, len(all))
+		copy(out, all)
+		return out
+	}
+	sIdx := sort.Search(len(all), func(i int) bool {
+		return opt.compareToPrefix(all[i].Biggest()) >= 0
+	})
+	if sIdx == len(all) {
+		// Not found.
+		return []*table.Table{}
+	}
+
+	filtered := all[sIdx:]
+	if !opt.prefixIsKey {
+		eIdx := sort.Search(len(filtered), func(i int) bool {
+			return opt.compareToPrefix(filtered[i].Smallest()) > 0
+		})
+		out := make([]*table.Table, len(filtered[:eIdx]))
+		copy(out, filtered[:eIdx])
+		return out
+	}
+
+	var out []*table.Table
+	for _, t := range filtered {
+		// When we encounter the first table whose smallest key is higher than
+		// opt.Prefix, we can stop.
+		if opt.compareToPrefix(t.Smallest()) > 0 {
+			return out
+		}
+		// opt.Prefix is actually the key. So, we can run bloom filter checks
+		// as well.
+		if t.DoesNotHave(opt.Prefix) {
+			continue
+		}
+		out = append(out, t)
+	}
+	return out
+}
+
 // DefaultIteratorOptions contains default options when iterating over Badger key-value stores.
 var DefaultIteratorOptions = IteratorOptions{
 	PrefetchValues: true,

iterator_test.go

@@ -27,6 +27,7 @@ import (
 	"testing"
 
 	"github.com/dgraph-io/badger/options"
+	"github.com/dgraph-io/badger/table"
 	"github.com/dgraph-io/badger/y"
 	"github.com/stretchr/testify/require"
 )
@@ -70,6 +71,59 @@ func TestPickTables(t *testing.T) {
 	outside("abd", "ab", "abc123")
 }
 
+func TestPickSortTables(t *testing.T) {
+	type MockKeys struct {
+		small string
+		large string
+	}
+	genTables := func(mks ...MockKeys) []*table.Table {
+		out := make([]*table.Table, 0)
+		for _, mk := range mks {
+			f := buildTable(t, [][]string{{mk.small, "some value"}, {mk.large, "some value"}})
+			opts := table.Options{LoadingMode: options.MemoryMap,
+				ChkMode: options.OnTableAndBlockRead}
+			tbl, err := table.OpenTable(f, opts)
+			require.NoError(t, err)
+			out = append(out, tbl)
+		}
+		return out
+	}
+	tables := genTables(MockKeys{small: "a", large: "abc"},
+		MockKeys{small: "abcd", large: "cde"},
+		MockKeys{small: "cge", large: "chf"},
+		MockKeys{small: "glr", large: "gyup"})
+	opt := DefaultIteratorOptions
+	opt.prefixIsKey = false
+	opt.Prefix = []byte("c")
+	filtered := opt.pickTables(tables)
+	require.Equal(t, 2, len(filtered))
+	// build table adds time stamp so removing tailing bytes.
+	require.Equal(t, filtered[0].Smallest()[:4], []byte("abcd"))
+	require.Equal(t, filtered[1].Smallest()[:3], []byte("cge"))
+	tables = genTables(MockKeys{small: "a", large: "abc"},
+		MockKeys{small: "abcd", large: "ade"},
+		MockKeys{small: "cge", large: "chf"},
+		MockKeys{small: "glr", large: "gyup"})
+	filtered = opt.pickTables(tables)
+	require.Equal(t, 1, len(filtered))
+	require.Equal(t, filtered[0].Smallest()[:3], []byte("cge"))
+	tables = genTables(MockKeys{small: "a", large: "abc"},
+		MockKeys{small: "abcd", large: "ade"},
+		MockKeys{small: "cge", large: "chf"},
+		MockKeys{small: "ckr", large: "cyup"},
+		MockKeys{small: "csfr", large: "gyup"})
+	filtered = opt.pickTables(tables)
+	require.Equal(t, 3, len(filtered))
+	require.Equal(t, filtered[0].Smallest()[:3], []byte("cge"))
+	require.Equal(t, filtered[1].Smallest()[:3], []byte("ckr"))
+	require.Equal(t, filtered[2].Smallest()[:4], []byte("csfr"))
+
+	opt.Prefix = []byte("aa")
+	filtered = opt.pickTables(tables)
+	require.Equal(t, y.ParseKey(filtered[0].Smallest()), []byte("a"))
+	require.Equal(t, y.ParseKey(filtered[0].Biggest()), []byte("abc"))
+}
+
 func TestIteratePrefix(t *testing.T) {
 	runBadgerTest(t, nil, func(t *testing.T, db *DB) {
 		bkey := func(i int) []byte {

level_handler.go

@@ -262,21 +262,23 @@ func (s *levelHandler) appendIterators(iters []y.Iterator, opt *IteratorOptions)
 	s.RLock()
 	defer s.RUnlock()
 
-	tables := make([]*table.Table, 0, len(s.tables))
-	for _, t := range s.tables {
-		if opt.pickTable(t) {
-			tables = append(tables, t)
+	if s.level == 0 {
+		// Remember to add in reverse order!
+		// The newer table at the end of s.tables should be added first as it takes precedence.
+		// Level 0 tables are not in key sorted order, so we need to consider them one by one.
+		var out []*table.Table
+		for _, t := range s.tables {
+			if opt.pickTable(t) {
+				out = append(out, t)
+			}
 		}
+		return appendIteratorsReversed(iters, out, opt.Reverse)
 	}
+
+	tables := opt.pickTables(s.tables)
 	if len(tables) == 0 {
 		return iters
 	}
-
-	if s.level == 0 {
-		// Remember to add in reverse order!
-		// The newer table at the end of s.tables should be added first as it takes precedence.
-		return appendIteratorsReversed(iters, tables, opt.Reverse)
-	}
 	return append(iters, table.NewConcatIterator(tables, opt.Reverse))
 }
 

table/iterator.go

@@ -402,7 +402,12 @@ type ConcatIterator struct {
 func NewConcatIterator(tbls []*Table, reversed bool) *ConcatIterator {
 	iters := make([]*Iterator, len(tbls))
 	for i := 0; i < len(tbls); i++ {
-		iters[i] = tbls[i].NewIterator(reversed)
+		// Increment the reference count. Since, we're not creating the iterator right now.
+		// Here, We'll hold the reference of the tables, till the lifecycle of the iterator.
+		tbls[i].IncrRef()
+
+		// Save cycles by not initializing the iterators until needed.
+		// iters[i] = tbls[i].NewIterator(reversed)
 	}
 	return &ConcatIterator{
 		reversed: reversed,
@@ -416,9 +421,12 @@ func (s *ConcatIterator) setIdx(idx int) {
 	s.idx = idx
 	if idx < 0 || idx >= len(s.iters) {
 		s.cur = nil
-	} else {
-		s.cur = s.iters[s.idx]
+		return
+	}
+	if s.iters[idx] == nil {
+		s.iters[idx] = s.tables[idx].NewIterator(s.reversed)
 	}
+	s.cur = s.iters[s.idx]
 }
 
 // Rewind implements y.Interface
@@ -498,7 +506,16 @@ func (s *ConcatIterator) Next() {
 
 // Close implements y.Interface.
 func (s *ConcatIterator) Close() error {
+	for _, t := range s.tables {
+		// DeReference the tables while closing the iterator.
+		if err := t.DecrRef(); err != nil {
+			return err
+		}
+	}
 	for _, it := range s.iters {
+		if it == nil {
+			continue
+		}
 		if err := it.Close(); err != nil {
 			return errors.Wrap(err, "ConcatIterator")
 		}