Skip to content

Commit

Permalink
Added AscendHint, DescendHint, and SeekHint to BTree and BTreeG
Browse files Browse the repository at this point in the history 
These new functions allow for using a path hint with iterators,
making it potentially faster when seeking to the first item in the iteration.

Benchmarks:

https://github.com/tidwall/btree-benchmark

About path hints:

https://github.com/tidwall/btree/blob/master/PATH_HINT.md

Usage:

```go
tr.AscendHint(key, iter, &hint)   // iterate items that are >= key, ascending
tr.DescendHint(key, iter, &hint)  // iterate items that are <= key, descending
iter.SeekHint(key, &hint)         // seek to item that is >= key
```
  • Loading branch information
@tidwall
tidwall committed Sep 6, 2023
1 parent 6b26091 commit 8d3586d
Show file tree
Hide file tree
Showing 4 changed files with 231 additions and 9 deletions.
6 changes: 6 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,9 @@ Load(item) // load presorted items into tree
SetHint(item, *hint) // insert or replace an existing item
GetHint(item, *hint) // get an existing item
DeleteHint(item, *hint) // delete an item
AscendHint(key, iter, *hint)
DescendHint(key, iter, *hint)
SeekHint(key, iter, *hint)

// Copy-on-write
Copy() // copy the btree
Expand Down Expand Up @@ -339,6 +342,9 @@ Load(item) // load presorted items into tree
SetHint(item, *hint) // insert or replace an existing item
GetHint(item, *hint) // get an existing item
DeleteHint(item, *hint) // delete an item
AscendHint(key, iter, *hint)
DescendHint(key, iter, *hint)
SeekHint(key, iter, *hint)

// Copy-on-write
Copy() // copy the btree
Expand Down
44 changes: 44 additions & 0 deletions btree.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,26 @@ func (tr *BTree) AscendMut(pivot any, iter func(item any) bool) {
}
}

func (tr *BTree) AscendHint(pivot any, iter func(item any) bool,
hint *PathHint,
) {
if pivot == nil {
tr.base.Scan(iter)
} else {
tr.base.AscendHint(pivot, iter, hint)
}
}

func (tr *BTree) AscendHintMut(pivot any, iter func(item any) bool,
hint *PathHint,
) {
if pivot == nil {
tr.base.ScanMut(iter)
} else {
tr.base.AscendHintMut(pivot, iter, hint)
}
}

// Descend the tree within the range [pivot, first]
// Pass nil for pivot to scan all item in descending order
// Return false to stop iterating
Expand All @@ -161,6 +181,26 @@ func (tr *BTree) DescendMut(pivot any, iter func(item any) bool) {
}
}

func (tr *BTree) DescendHint(pivot any, iter func(item any) bool,
hint *PathHint,
) {
if pivot == nil {
tr.base.Reverse(iter)
} else {
tr.base.DescendHint(pivot, iter, hint)
}
}

func (tr *BTree) DescendHintMut(pivot any, iter func(item any) bool,
hint *PathHint,
) {
if pivot == nil {
tr.base.ReverseMut(iter)
} else {
tr.base.DescendHintMut(pivot, iter, hint)
}
}

// Load is for bulk loading pre-sorted items
// If the load replaces and existing item then the value for the replaced item
// is returned.
Expand Down Expand Up @@ -317,6 +357,10 @@ func (iter *Iter) Seek(key any) bool {
return iter.base.Seek(key)
}

func (iter *Iter) SeekHint(key any, hint *PathHint) bool {
return iter.base.SeekHint(key, hint)
}

// First moves iterator to first item in tree.
// Returns false if the tree is empty.
func (iter *Iter) First() bool {
Expand Down
54 changes: 45 additions & 9 deletions btreeg.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -601,19 +601,30 @@ func (tr *BTreeG[T]) nodeRebalance(n *node[T], i int) {
// Pass nil for pivot to scan all item in ascending order
// Return false to stop iterating
func (tr *BTreeG[T]) Ascend(pivot T, iter func(item T) bool) {
tr.ascend(pivot, iter, false)
tr.ascend(pivot, iter, false, nil)
}
func (tr *BTreeG[T]) AscendMut(pivot T, iter func(item T) bool) {
tr.ascend(pivot, iter, true)
tr.ascend(pivot, iter, true, nil)
}
func (tr *BTreeG[T]) ascend(pivot T, iter func(item T) bool, mut bool) {
func (tr *BTreeG[T]) ascend(pivot T, iter func(item T) bool, mut bool,
hint *PathHint,
) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return
}
tr.nodeAscend(&tr.root, pivot, nil, 0, iter, mut)
tr.nodeAscend(&tr.root, pivot, hint, 0, iter, mut)
}
func (tr *BTreeG[T]) AscendHint(pivot T, iter func(item T) bool, hint *PathHint,
) {
tr.ascend(pivot, iter, false, hint)
}
func (tr *BTreeG[T]) AscendHintMut(pivot T, iter func(item T) bool,
hint *PathHint,
) {
tr.ascend(pivot, iter, true, hint)
}

// The return value of this function determines whether we should keep iterating
Expand Down Expand Up @@ -693,19 +704,32 @@ func (tr *BTreeG[T]) nodeReverse(cn **node[T], iter func(item T) bool, mut bool,
// Pass nil for pivot to scan all item in descending order
// Return false to stop iterating
func (tr *BTreeG[T]) Descend(pivot T, iter func(item T) bool) {
tr.descend(pivot, iter, false)
tr.descend(pivot, iter, false, nil)
}
func (tr *BTreeG[T]) DescendMut(pivot T, iter func(item T) bool) {
tr.descend(pivot, iter, true)
tr.descend(pivot, iter, true, nil)
}
func (tr *BTreeG[T]) descend(pivot T, iter func(item T) bool, mut bool) {
func (tr *BTreeG[T]) descend(pivot T, iter func(item T) bool, mut bool,
hint *PathHint,
) {
if tr.lock(mut) {
defer tr.unlock(mut)
}
if tr.root == nil {
return
}
tr.nodeDescend(&tr.root, pivot, nil, 0, iter, mut)
tr.nodeDescend(&tr.root, pivot, hint, 0, iter, mut)
}

func (tr *BTreeG[T]) DescendHint(pivot T, iter func(item T) bool,
hint *PathHint,
) {
tr.descend(pivot, iter, false, hint)
}
func (tr *BTreeG[T]) DescendHintMut(pivot T, iter func(item T) bool,
hint *PathHint,
) {
tr.descend(pivot, iter, true, hint)
}

func (tr *BTreeG[T]) nodeDescend(cn **node[T], pivot T, hint *PathHint,
Expand Down Expand Up @@ -1104,6 +1128,7 @@ type IterG[T any] struct {
seeked bool
atstart bool
atend bool
stack0 [4]iterStackItemG[T]
stack []iterStackItemG[T]
item T
}
Expand All @@ -1128,12 +1153,21 @@ func (tr *BTreeG[T]) iter(mut bool) IterG[T] {
iter.tr = tr
iter.mut = mut
iter.locked = tr.lock(iter.mut)
iter.stack = iter.stack0[:0]
return iter
}

// Seek to item greater-or-equal-to key.
// Returns false if there was no item found.
func (iter *IterG[T]) Seek(key T) bool {
return iter.seek(key, nil)
}

func (iter *IterG[T]) SeekHint(key T, hint *PathHint) bool {
return iter.seek(key, hint)
}

func (iter *IterG[T]) seek(key T, hint *PathHint) bool {
if iter.tr == nil {
return false
}
Expand All @@ -1143,8 +1177,9 @@ func (iter *IterG[T]) Seek(key T) bool {
return false
}
n := iter.tr.isoLoad(&iter.tr.root, iter.mut)
var depth int
for {
i, found := iter.tr.find(n, key, nil, 0)
i, found := iter.tr.find(n, key, hint, depth)
iter.stack = append(iter.stack, iterStackItemG[T]{n, i})
if found {
iter.item = n.items[i]
Expand All @@ -1155,6 +1190,7 @@ func (iter *IterG[T]) Seek(key T) bool {
return iter.Next()
}
n = iter.tr.isoLoad(&(*n.children)[i], iter.mut)
depth++
}
}

Expand Down
136 changes: 136 additions & 0 deletions btreeg_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,66 @@ func TestGenericDescend(t *testing.T) {
}
}

func TestGenericDescendHint(t *testing.T) {
tr := testNewBTree()
var count int
var hint PathHint
tr.DescendHint(testMakeItem(rand.Int()), func(item testKind) bool {
count++
return true
}, &hint)
if count > 0 {
t.Fatalf("expected 0, got %v", count)
}
var keys []testKind
for i := 0; i < 1000; i += 10 {
keys = append(keys, testMakeItem(i))
tr.Set(keys[len(keys)-1])
}
var exp []testKind
tr.Reverse(func(item testKind) bool {
exp = append(exp, item)
return true
})
for i := 999; i >= 0; i-- {
key := testMakeItem(i)
var all []testKind
tr.DescendHint(key, func(item testKind) bool {
all = append(all, item)
return true
}, &hint)
for len(exp) > 0 && tr.Less(key, exp[0]) {
exp = exp[1:]
}
var count int
tr.DescendHint(key, func(item testKind) bool {
if count == (i+1)%tr.max {
return false
}
count++
return true
}, &hint)
if count > len(exp) {
t.Fatalf("expected 1, got %v", count)
}
if !kindsAreEqual(exp, all) {
fmt.Printf("exp: %v\n", exp)
fmt.Printf("all: %v\n", all)
t.Fatal("mismatch")
}
for j := 0; j < tr.Len(); j++ {
count = 0
tr.DescendHint(key, func(item testKind) bool {
if count == j {
return false
}
count++
return true
}, &hint)
}
}
}

func TestGenericAscend(t *testing.T) {
tr := testNewBTree()
var count int
Expand Down Expand Up @@ -190,6 +250,51 @@ func TestGenericAscend(t *testing.T) {
}
}

func TestGenericAscendHint(t *testing.T) {
tr := testNewBTree()
var hint PathHint
var count int
tr.AscendHint(testMakeItem(1), func(item testKind) bool {
count++
return true
}, &hint)
if count > 0 {
t.Fatalf("expected 0, got %v", count)
}
var keys []testKind
for i := 0; i < 1000; i += 10 {
keys = append(keys, testMakeItem(i))
tr.Set(keys[len(keys)-1])
tr.sane()
}
exp := keys
for i := -1; i < 1000; i++ {
key := testMakeItem(i)
var all []testKind
tr.AscendHint(key, func(item testKind) bool {
all = append(all, item)
return true
}, &hint)
for len(exp) > 0 && tr.Less(exp[0], key) {
exp = exp[1:]
}
var count int
tr.AscendHint(key, func(item testKind) bool {
if count == (i+1)%tr.max {
return false
}
count++
return true
}, &hint)
if count > len(exp) {
t.Fatalf("expected 1, got %v", count)
}
if !kindsAreEqual(exp, all) {
t.Fatal("mismatch")
}
}
}

func TestGenericItems(t *testing.T) {
tr := testNewBTree()
if len(tr.Items()) != 0 {
Expand Down Expand Up @@ -1328,6 +1433,37 @@ func TestGenericIterSeek(t *testing.T) {
}
}

func TestGenericIterSeekHint(t *testing.T) {
tr := NewBTreeG(func(a, b int) bool {
return a < b
})
var all []int
for i := 0; i < 10000; i++ {
tr.Set(i * 2)
all = append(all, i)
}
var hint PathHint
_ = all
{
iter := tr.Iter()
var vals []int
for ok := iter.SeekHint(501, &hint); ok; ok = iter.Next() {
vals = append(vals, iter.Item())
}
iter.Release()
assert(vals[0] == 502 && vals[1] == 504)
}
{
iter := tr.Iter()
var vals []int
for ok := iter.SeekHint(501, &hint); ok; ok = iter.Prev() {
vals = append(vals, iter.Item())
}
iter.Release()
assert(vals[0] == 502 && vals[1] == 500)
}
}

func TestGenericIterSeekPrefix(t *testing.T) {
tr := NewBTreeG(func(a, b int) bool {
return a < b
Expand Down

0 comments on commit 8d3586d

Please sign in to comment.