B+ Tree — Hands-On Tasks¶

package bplustree

const T = 3 // minimum number of keys per non-root node

type Node interface {
    isLeaf() bool
}

type Leaf struct {
    keys []int
    vals []string
    next *Leaf
}

func (*Leaf) isLeaf() bool { return true }

type Internal struct {
    keys     []int
    children []Node
}

func (*Internal) isLeaf() bool { return false }

type Tree struct {
    root Node
}

func New() *Tree { return &Tree{root: &Leaf{}} }

func (t *Tree) Search(k int) (string, bool) {
    n := t.root
    for !n.isLeaf() {
        in := n.(*Internal)
        i := upperBound(in.keys, k)
        n = in.children[i]
    }
    lf := n.(*Leaf)
    i := lowerBound(lf.keys, k)
    if i < len(lf.keys) && lf.keys[i] == k {
        return lf.vals[i], true
    }
    return "", false
}

func (t *Tree) Insert(k int, v string) {
    sepKey, newChild, split := insert(t.root, k, v)
    if split {
        t.root = &Internal{
            keys:     []int{sepKey},
            children: []Node{t.root, newChild},
        }
    }
}

func insert(n Node, k int, v string) (int, Node, bool) {
    if n.isLeaf() {
        lf := n.(*Leaf)
        i := lowerBound(lf.keys, k)
        if i < len(lf.keys) && lf.keys[i] == k {
            lf.vals[i] = v
            return 0, nil, false
        }
        lf.keys = insertAt(lf.keys, i, k)
        lf.vals = insertStrAt(lf.vals, i, v)
        if len(lf.keys) <= 2*T {
            return 0, nil, false
        }
        mid := len(lf.keys) / 2
        right := &Leaf{
            keys: append([]int{}, lf.keys[mid:]...),
            vals: append([]string{}, lf.vals[mid:]...),
            next: lf.next,
        }
        lf.keys = lf.keys[:mid]
        lf.vals = lf.vals[:mid]
        lf.next = right
        return right.keys[0], right, true
    }
    in := n.(*Internal)
    i := upperBound(in.keys, k)
    sep, child, split := insert(in.children[i], k, v)
    if !split {
        return 0, nil, false
    }
    in.keys = insertAt(in.keys, i, sep)
    in.children = insertChildAt(in.children, i+1, child)
    if len(in.keys) <= 2*T {
        return 0, nil, false
    }
    mid := len(in.keys) / 2
    sepUp := in.keys[mid]
    right := &Internal{
        keys:     append([]int{}, in.keys[mid+1:]...),
        children: append([]Node{}, in.children[mid+1:]...),
    }
    in.keys = in.keys[:mid]
    in.children = in.children[:mid+1]
    return sepUp, right, true
}

// Delete, borrow, merge, lowerBound, upperBound, and insertAt helpers
// follow classic B-tree mechanics but mind the leaf-chain pointer:
// on a leaf merge, splice the merged leaf's `next` into the predecessor's `next`.

func (t *Tree) RangeScan(lo, hi int) []int {
    n := t.root
    for !n.isLeaf() {
        in := n.(*Internal)
        i := upperBound(in.keys, lo)
        n = in.children[i]
    }
    var out []int
    lf := n.(*Leaf)
    i := lowerBound(lf.keys, lo)
    for lf != nil {
        for ; i < len(lf.keys); i++ {
            if lf.keys[i] > hi {
                return out
            }
            out = append(out, lf.keys[i])
        }
        lf = lf.next
        i = 0
    }
    return out
}

type Internal struct {
    keys      []int
    children  []Node
    highKey   int
    rightLink *Internal
}
type Leaf struct {
    keys      []int
    vals      []string
    highKey   int
    next      *Leaf // also serves as rightLink
}

func (t *Tree) ConcurrentSearch(k int) (string, bool) {
    n := t.root
    for !n.isLeaf() {
        in := n.(*Internal)
        for k > in.highKey && in.rightLink != nil {
            in = in.rightLink
        }
        i := upperBound(in.keys, k)
        n = in.children[i]
    }
    lf := n.(*Leaf)
    for k > lf.highKey && lf.next != nil {
        lf = lf.next
    }
    i := lowerBound(lf.keys, k)
    if i < len(lf.keys) && lf.keys[i] == k {
        return lf.vals[i], true
    }
    return "", false
}

func BulkLoad(pairs []struct{ K int; V string }, fillRatio float64) *Tree {
    leafCap := int(float64(2*T) * fillRatio)
    if leafCap < T { leafCap = T }

    var leaves []*Leaf
    cur := &Leaf{}
    for _, p := range pairs {
        if len(cur.keys) == leafCap {
            leaves = append(leaves, cur)
            nxt := &Leaf{}
            cur.next = nxt
            cur = nxt
        }
        cur.keys = append(cur.keys, p.K)
        cur.vals = append(cur.vals, p.V)
    }
    if len(cur.keys) > 0 {
        leaves = append(leaves, cur)
    }
    if len(leaves) == 1 {
        return &Tree{root: leaves[0]}
    }
    var level []Node
    for _, lf := range leaves { level = append(level, lf) }
    seps := make([]int, 0, len(leaves)-1)
    for i := 1; i < len(leaves); i++ {
        seps = append(seps, leaves[i].keys[0])
    }
    for len(level) > 1 {
        nextLevel := []Node{}
        nextSeps := []int{}
        intCap := 2 * T
        i := 0
        for i < len(level) {
            in := &Internal{}
            end := i + intCap + 1
            if end > len(level) { end = len(level) }
            in.children = level[i:end]
            if i > 0 {
                nextSeps = append(nextSeps, seps[i-1])
            }
            for j := i; j < end-1; j++ {
                in.keys = append(in.keys, seps[j])
            }
            nextLevel = append(nextLevel, in)
            i = end
        }
        level = nextLevel
        seps = nextSeps
    }
    return &Tree{root: level[0]}
}

type CompressedLeaf struct {
    prefix []byte
    suffixes [][]byte
    vals []string
    next *CompressedLeaf
}

func (lf *CompressedLeaf) full(i int) []byte {
    b := make([]byte, 0, len(lf.prefix)+len(lf.suffixes[i]))
    b = append(b, lf.prefix...)
    b = append(b, lf.suffixes[i]...)
    return b
}

func commonPrefix(a, b []byte) int {
    n := len(a)
    if len(b) < n { n = len(b) }
    for i := 0; i < n; i++ {
        if a[i] != b[i] { return i }
    }
    return n
}

func compressLeaf(keys [][]byte, vals []string) *CompressedLeaf {
    if len(keys) == 0 { return &CompressedLeaf{} }
    p := commonPrefix(keys[0], keys[len(keys)-1])
    out := &CompressedLeaf{prefix: keys[0][:p], vals: vals}
    for _, k := range keys {
        out.suffixes = append(out.suffixes, append([]byte{}, k[p:]...))
    }
    return out
}

const PageSize = 4096
const (
    PageLeaf     byte = 1
    PageInternal byte = 2
)

type Pager struct {
    f    *os.File
    next uint32
    lru  map[uint32]*Page
}

type Page struct {
    id    uint32
    dirty bool
    data  [PageSize]byte
}

func (p *Pager) Read(id uint32) (*Page, error) {
    if pg, ok := p.lru[id]; ok { return pg, nil }
    pg := &Page{id: id}
    _, err := p.f.ReadAt(pg.data[:], int64(id)*PageSize)
    if err != nil { return nil, err }
    p.lru[id] = pg
    return pg, nil
}

func (p *Pager) Write(pg *Page) error {
    _, err := p.f.WriteAt(pg.data[:], int64(pg.id)*PageSize)
    pg.dirty = false
    return err
}

// Leaf encoding: [type:1][n:2][nextLeaf:4][ keys[]:4 * n ][ valOffs[]:2 * n ][ vals[] ]
// Fits ~340 int4 keys + int4 row pointers per 4 KB leaf depending on value size.

type CowTree struct {
    roots []Node // roots[i] is the root at snapshot i
}

func cowInsert(n Node, k int, v string) (Node, int, Node, bool) {
    if n.isLeaf() {
        old := n.(*Leaf)
        nl := &Leaf{
            keys: append([]int{}, old.keys...),
            vals: append([]string{}, old.vals...),
            next: old.next,
        }
        i := lowerBound(nl.keys, k)
        if i < len(nl.keys) && nl.keys[i] == k {
            nl.vals[i] = v
        } else {
            nl.keys = insertAt(nl.keys, i, k)
            nl.vals = insertStrAt(nl.vals, i, v)
        }
        if len(nl.keys) <= 2*T {
            return nl, 0, nil, false
        }
        mid := len(nl.keys) / 2
        right := &Leaf{
            keys: append([]int{}, nl.keys[mid:]...),
            vals: append([]string{}, nl.vals[mid:]...),
            next: nl.next,
        }
        nl.keys = nl.keys[:mid]
        nl.vals = nl.vals[:mid]
        nl.next = right
        return nl, right.keys[0], right, true
    }
    old := n.(*Internal)
    ni := &Internal{
        keys:     append([]int{}, old.keys...),
        children: append([]Node{}, old.children...),
    }
    i := upperBound(ni.keys, k)
    newChild, sep, sib, split := cowInsert(ni.children[i], k, v)
    ni.children[i] = newChild
    if !split {
        return ni, 0, nil, false
    }
    ni.keys = insertAt(ni.keys, i, sep)
    ni.children = insertChildAt(ni.children, i+1, sib)
    if len(ni.keys) <= 2*T {
        return ni, 0, nil, false
    }
    mid := len(ni.keys) / 2
    sepUp := ni.keys[mid]
    right := &Internal{
        keys:     append([]int{}, ni.keys[mid+1:]...),
        children: append([]Node{}, ni.children[mid+1:]...),
    }
    ni.keys = ni.keys[:mid]
    ni.children = ni.children[:mid+1]
    return ni, sepUp, right, true
}

func (t *CowTree) Insert(k int, v string) int {
    root := t.roots[len(t.roots)-1]
    newRoot, sep, sib, split := cowInsert(root, k, v)
    if split {
        newRoot = &Internal{keys: []int{sep}, children: []Node{newRoot, sib}}
    }
    t.roots = append(t.roots, newRoot)
    return len(t.roots) - 1
}