From fa7bb702a7d041b2df6d40b867ce02ed74505d50 Mon Sep 17 00:00:00 2001 From: NamanJain8 Date: Tue, 10 Nov 2020 13:47:24 +0530 Subject: [PATCH 1/4] add btree reset option and btree with file --- z/btree.go | 34 +++++++++++++++++++++++++++++----- z/btree_test.go | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/z/btree.go b/z/btree.go index 28a1c136..50aae2a0 100644 --- a/z/btree.go +++ b/z/btree.go @@ -48,16 +48,27 @@ func (t *Tree) Release() { } } -// NewTree returns a memory mapped B+ tree. +func createFile(maxSz int, fname string) (*MmapFile, error) { + if fname == "" { + fd, err := ioutil.TempFile("", "btree") + check(err) + return OpenMmapFileUsing(fd, maxSz, true) + } + return OpenMmapFile(fname, os.O_RDWR|os.O_CREATE, maxSz) +} + +// NewTree returns a new memory mapped tree in a tmp directory. func NewTree(maxSz int) *Tree { - // Tell kernel that we'd be reading pages in random order, so don't do read ahead. - fd, err := ioutil.TempFile("", "btree") - check(err) + return NewTreeWithFile(maxSz, "") +} - mf, err := OpenMmapFileUsing(fd, maxSz, true) +// NewTreeWithFile returns a memory mapped B+ tree with given filename. +func NewTreeWithFile(maxSz int, fname string) *Tree { + mf, err := createFile(maxSz, fname) if err != NewFile { check(err) } + // Tell kernel that we'd be reading pages in random order, so don't do read ahead. check(Madvise(mf.Data, false)) t := &Tree{ @@ -72,6 +83,19 @@ func NewTree(maxSz int) *Tree { return t } +// Reset resets the tree and truncates it to maxSz. +func (t *Tree) Reset(maxSz int) { + t.nextPage = 1 + t.freePage = 0 + if maxSz < pageSize { + maxSz = pageSize + } + check(t.mf.Truncate(int64(maxSz))) + // Set the root. + t.newNode(0) + t.Set(absoluteMax, 0) +} + type TreeStats struct { NextPage int NumNodes int diff --git a/z/btree_test.go b/z/btree_test.go index 3467459c..aa7047ec 100644 --- a/z/btree_test.go +++ b/z/btree_test.go @@ -82,6 +82,39 @@ func TestTreeBasic(t *testing.T) { setAndGet() } +func TestTreeReset(t *testing.T) { + bt := NewTreeWithFile(1<<20, "/home/algod/btree") + N := 1 << 10 + val := rand.Uint64() + for i := 0; i < N; i++ { + bt.Set(rand.Uint64(), val) + } + + // Truncate it to small size that is less than pageSize. + bt.Reset(1 << 10) + + stats := bt.Stats() + // Verify the tree stats. + require.Equal(t, 3, stats.NextPage) + require.Equal(t, 2, stats.NumNodes) + require.Equal(t, 1, stats.NumLeafKeys) + require.Equal(t, 1, stats.NumLeafKeys) + require.Equal(t, 2*pageSize, stats.Bytes) + expectedOcc := float64(2) * 100 / float64(stats.NumNodes*maxKeys) + require.InDelta(t, expectedOcc, stats.Occupancy, 0.01) + require.Zero(t, stats.FreePages) + // Check if we can reinsert the data. + mp := make(map[uint64]uint64) + for i := 0; i < N; i++ { + k := rand.Uint64() + mp[k] = val + bt.Set(k, val) + } + for k, v := range mp { + require.Equal(t, v, bt.Get(k)) + } +} + func TestTreeCycle(t *testing.T) { bt := NewTree(1 << 20) val := uint64(0) From 8f1b1734d53ea4e00234a8d7dc185d82c5a81b82 Mon Sep 17 00:00:00 2001 From: NamanJain8 Date: Tue, 10 Nov 2020 13:51:06 +0530 Subject: [PATCH 2/4] clean --- z/btree_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/z/btree_test.go b/z/btree_test.go index aa7047ec..4c2cb6e6 100644 --- a/z/btree_test.go +++ b/z/btree_test.go @@ -83,7 +83,7 @@ func TestTreeBasic(t *testing.T) { } func TestTreeReset(t *testing.T) { - bt := NewTreeWithFile(1<<20, "/home/algod/btree") + bt := NewTree(1 << 20) N := 1 << 10 val := rand.Uint64() for i := 0; i < N; i++ { From f777425d4ddbb12611b1430aacc9cab838358ee3 Mon Sep 17 00:00:00 2001 From: NamanJain8 Date: Tue, 10 Nov 2020 14:55:14 +0530 Subject: [PATCH 3/4] clean API --- z/btree.go | 9 ++------- z/btree_test.go | 18 +++++++++--------- 2 files changed, 11 insertions(+), 16 deletions(-) diff --git a/z/btree.go b/z/btree.go index 50aae2a0..a0786457 100644 --- a/z/btree.go +++ b/z/btree.go @@ -57,13 +57,8 @@ func createFile(maxSz int, fname string) (*MmapFile, error) { return OpenMmapFile(fname, os.O_RDWR|os.O_CREATE, maxSz) } -// NewTree returns a new memory mapped tree in a tmp directory. -func NewTree(maxSz int) *Tree { - return NewTreeWithFile(maxSz, "") -} - -// NewTreeWithFile returns a memory mapped B+ tree with given filename. -func NewTreeWithFile(maxSz int, fname string) *Tree { +// NewTree returns a memory mapped B+ tree with given filename. +func NewTree(maxSz int, fname string) *Tree { mf, err := createFile(maxSz, fname) if err != NewFile { check(err) diff --git a/z/btree_test.go b/z/btree_test.go index 4c2cb6e6..1e4c055b 100644 --- a/z/btree_test.go +++ b/z/btree_test.go @@ -36,8 +36,7 @@ func setPageSize(sz int) { } func TestTree(t *testing.T) { - bt := NewTree(1 << 20) - // bt.Print() + bt := NewTree(1<<20, "") N := uint64(256 * 256) for i := uint64(1); i < N; i++ { @@ -54,12 +53,11 @@ func TestTree(t *testing.T) { for i := uint64(100); i < N; i++ { require.Equal(t, i, bt.Get(i)) } - // bt.Print() } func TestTreeBasic(t *testing.T) { setAndGet := func() { - bt := NewTree(1 << 20) + bt := NewTree(1<<20, "") defer bt.Release() N := uint64(1 << 20) @@ -83,7 +81,8 @@ func TestTreeBasic(t *testing.T) { } func TestTreeReset(t *testing.T) { - bt := NewTree(1 << 20) + bt := NewTree(1<<20, "") + defer bt.Release() N := 1 << 10 val := rand.Uint64() for i := 0; i < N; i++ { @@ -116,7 +115,8 @@ func TestTreeReset(t *testing.T) { } func TestTreeCycle(t *testing.T) { - bt := NewTree(1 << 20) + bt := NewTree(1<<20, "") + defer bt.Release() val := uint64(0) for i := 0; i < 16; i++ { for j := 0; j < 1e6+i*1e4; j++ { @@ -142,7 +142,7 @@ func TestOccupancyRatio(t *testing.T) { defer setPageSize(os.Getpagesize()) require.Equal(t, 4, maxKeys) - bt := NewTree(1 << 20) + bt := NewTree(1<<20, "") defer bt.Release() expectedRatio := float64(1) * 100 / float64(maxKeys) @@ -245,7 +245,7 @@ func BenchmarkWrite(b *testing.B) { } }) b.Run("btree", func(b *testing.B) { - bt := NewTree(1 << 30) + bt := NewTree(1<<30, "") defer bt.Release() b.ResetTimer() for n := 0; n < b.N; n++ { @@ -279,7 +279,7 @@ func BenchmarkRead(b *testing.B) { } }) - bt := NewTree(1 << 30) + bt := NewTree(1<<30, "") defer bt.Release() for i := 0; i < N; i++ { k := uint64(rand.Intn(2*N)) + 1 From dc61f31cdf323b8a1d276ed6098fe0b80758d717 Mon Sep 17 00:00:00 2001 From: NamanJain8 Date: Tue, 10 Nov 2020 17:03:51 +0530 Subject: [PATCH 4/4] initRootNode review comment --- z/btree.go | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/z/btree.go b/z/btree.go index a0786457..b9300cd0 100644 --- a/z/btree.go +++ b/z/btree.go @@ -57,6 +57,13 @@ func createFile(maxSz int, fname string) (*MmapFile, error) { return OpenMmapFile(fname, os.O_RDWR|os.O_CREATE, maxSz) } +func (t *Tree) initRootNode() { + // This is the root node. + t.newNode(0) + // This acts as the rightmost pointer (all the keys are <= this key). + t.Set(absoluteMax, 0) +} + // NewTree returns a memory mapped B+ tree with given filename. func NewTree(maxSz int, fname string) *Tree { mf, err := createFile(maxSz, fname) @@ -70,11 +77,7 @@ func NewTree(maxSz int, fname string) *Tree { mf: mf, nextPage: 1, } - // This is the root node. - t.newNode(0) - - // This acts as the rightmost pointer (all the keys are <= this key). - t.Set(absoluteMax, 0) + t.initRootNode() return t } @@ -86,9 +89,7 @@ func (t *Tree) Reset(maxSz int) { maxSz = pageSize } check(t.mf.Truncate(int64(maxSz))) - // Set the root. - t.newNode(0) - t.Set(absoluteMax, 0) + t.initRootNode() } type TreeStats struct {