Build Your Own Database From Scratch~04 BTree

The Practice

本章将在 Golang 中实现一棵不可变的 B+ 树。实现过程非常简单，因此很容易理解。

The Node Format

我们的B树最终会持久化到磁盘上，所以我们需要先设计B树节点的传输格式。如果没有格式，我们将不知道节点的大小以及何时拆分节点。

一个 node 包含：

| type | nkeys |  pointers  |   offsets  | key-values
|  2B  |   2B  | nkeys * 8B | nkeys * 2B | ...

• type：代表节点类型，这里为了简单只有两种节点类型（内节点和叶子节点）
• nkeys：代表页中一共有多少个 key
• pointers：存储了所有 key 所在页的地址
• offsets：存储了所有 key 在当前页中的偏移量
• key-values：实际记录（仅在叶子节点生效）

这就是 KV 对的格式，长度后跟数据。

| klen | vlen | key | val |
|  2B  |  2B  | ... | ... |

为了简单起见，叶节点和内部节点都使用相同的格式。

Data Types

既然我们最终要将 B 树转储到磁盘，为什么不使用字节数组作为我们的内存数据结构呢？

type BNode struct {
    data []byte  // can be dumped to the disk
}
 
const (
    BNODE_NODE = 1 // internal nodes without values 内节点
    BNODE_LEAF = 2 // leaf nodes with values        叶子节点
)

而且我们不能使用内存中的指针，指针是引用磁盘页的64位整数（也叫内存地址），而不是内存中的节点，我们将添加一些回调来抽象这个方面，这样我们的数据结构代码仍然是纯数据结构代码。

tpye BTree struct {
    // pointer ( a nonzero page number)
    root uint64
    // callbacks for managing on-disk pages
    get func(uint64) BNode // dereference a pointer
    new func(BNode) uint64 // allocate a new page
    del func(uint64)      // deallocate a page
}

页面大小定义为 4K 字节，更大的页面大小（如 8K 或 16K）也可以使用。

一个 uint64 的指针值唯一指向一个 BNode，即一个页

我们还添加了一些对键和值的大小的约束。因此，具有单个 KV 对的节点始终可以放在单个页面上。如果您需要支持更大的键或更大的值，则必须为它们分配额外的页面，这会增加复杂性。

const HEADER = 4
 
const BTREE_PAGE_SIZE = 4096
const BTREE_MAX_KEY_SIZE = 1000
const BTREE_MAX_VAL_SIZE = 3000
 
func init() {
    node1max := HEADER + 8 + 2 + 4 + BTREE_MAX_KEY_SIZE + BTREE_MAX_VAL_SIZE 
    assert(node1max <= BTREE_PAGE_SIZE)
}

Decoding the B-tree Node

由于节点只是一个字节数组，因此我们将添加一些辅助函数来访问其内容。

// header
func (node BNode) btype() uint16 {
    return binary.LittleEndian.Uint16(node.data)
}
 
func (node BNode) nkeys() uint16 {
    return binary.LittleEndian.Uint16(node.data[2:4])
}
 
func (nonde BNode) setHeader(btype uint16, nkeys uint16) {
    binary.LittleEndian.PutUint16(node.data[0:2], btype)
    binary.LittleEndian.PutUint16(node.data[2:4], nkeys)
}

// pointers
func (node BNode) getPtr(idx uint16) uint64 {
    assert(idx < node.nkeys())
    // 获取 idx 对应页的地址偏移量
    pos := HEADER + 8 * idx
    return binary.LittleEndian.Uint64(node.data[pos:])
}
 
func (node BNode) setPtr(idx uint64, val uint64) {
    assert(idx < node.nkeys())
    pos := HEADER + idx * 8
    binary.LittleEndian.PutUint64(node.data[pos:], val)
}

有关偏移列表的一些详细信息：

• 该偏移量是相对于第一个 KV 对的位置而言的。
• 第一个 KV 对的偏移量始终为零，因此不存储在列表中。
• 我们在偏移量列表中存储最后一个 KV 对末尾的偏移量，用来确定节点的大小。

// offset list
func offsetPos(node BNode, idx uint16) uint16 {
    assert(1 <= idx && idx <= node.nkeys())
    // idx 在 offset list 中的偏移量
    return HEADER  + 8 * node.nkeys() + 2 * (idx-1)
}
 
func (node BNode) getOffset(idx uint16) uint16 {
    if idx == 0 {
        return 0
    }
    // 获取 idx 对应节点的 kv 偏移量信息
    // 首先获取 idx 对应节点的在偏移量列表中的偏移量
    // 然后读取 2B 字节后返回 idx 对应的 KV 在数据部分的偏移量
    return binary.LittleEndian.Uint16(node.data[offsetPos(node, idx):])
}

// key-values
func (node BNode) kvPos(idx uint16) uint16 {
    assert(idx <= node.nkeys())
    // 类似与 golang map 定位 kv
    // HEADER + 8 * node.nkeys() + 2 * node.nkeys() 这里是记录元信息部分
    // 记录元信息部分之后全部都是数据部分，而 getOffset 已经求出了 kv 在这部分的偏移量，把记录元信息部分长度+kv 在数据部分的偏移量就得到了：idx 这个 kv 对相比于整个页的起始位置的偏移量
    return HEADER + 8 * node.nkeys() + 2 * node.nkeys() + node.getOffset(idx)
}
 
func (node BNode) getKey(idx uint16) []byte {
    assert(idx < node.nkeys())
    // idx 对应 kv 对的起始偏移量（相对于数据部分起始位置来说）
    pos := node.kvPos(idx)
    // klen(2B) + vlen(2B) + key + value
    klen := binary.LittleEndian.Uint16(node.data[pos:])
    // node.data[pos+4:] （第一个 kv 的数据部分起始位置）读取 klen 长度
    return node.data[pos+4:][:klen]
}
 
func (node BNode) getVal(idx uint16) []byte {
    assert(idx < node.nkeys())
    pos := node.kvPos(idx)
    // klen(2B) + vlen(2B) + key + value
    klen := binary.LittleEndian.Uint16(node.data[pos+0:])
    vlen := binary.LittleEndian.Uint16(node.data[pos+2:])
    return node.data[pos+4+klen:][:vlen]
}

并确定节点的大小。

// node size in bytes
func (node BNode) nbytes uint16 {
    // 调用 kvPos 的 idx = nkeys ，这相当于求最后一条 kv 相比于页开头的偏移量，那实际上就是求整个页的大小
    return node.kvPos(node.nkeys())
}

The B-Tree Insertion

该代码被分解为小步骤。

Step 1: Look Up the Key

要将键插入叶节点，我们需要在排序的 KV 列表中查找它的位置。

// returns the first kid node whose range intersects the key. (kid[i] <= key)
func nodeLookupLE(node BNode, key []byte) uint16 {
    nkeys := node.nkeys()
    found := uint16(0)
    
    // the first key is a copy from the parent node
    // thus it's always less than or equal to the key. 第一个 key 是父节点（根节点）的副本，因此它等于或者小于所有 key，这里不在插入逻辑中判断
    for i := uint16(1); i < nkeys; i++ {
        cmp := bytes.Compare(node.getKey(i), key)
        if cmp <= 0 { // 在该页找到比 key 小的 key 列表中最大的 key
            found = i
        }
        if cmp >= 0 { // 因为 key 是从小到大排列的，所以一旦找到一个 key >= target，那么之后所有的 key 都大于 target，因此可以直接退出循环了
            break
        }
    }
    return found
}

查找对叶节点和内部节点都有效。请注意，第一个键将被跳过比较，因为它已经与父节点比较过了。

Step 2: Update Leaf Nodes

查找到要插入的位置后，我们需要创建一个包含新键的节点副本。

// add a new key to a leaf node.
func leafInsert (new BNode, old BNode, idx uint16, key []byte, val []byte) {
    // 为新节点设置 header
    new.setHeader(BNODE_LEAF, old.nkeys()+1)
    // 插入前半部分 KV 对
    nodeAppendRange(new, old, 0, 0, idx)
    // 插入新的 KV 对
    nodeAppendKV(new, idx, 0, key, val)
    // 插入后半部分 KV 对
    nodeAppendRange(new, old, idx+1, idx, old.nkeys()-idx)
}

nodeAppendRange 函数将旧节点中的键复制到新节点中。

// copy multiple KVs into the position
func nodeAppendRange(new, old BNode, dstNew uint64, srcOld uint16, n uint16) {
    // 越界判断
    assert(srcOld+n <= old.nkeys())
    assert(dstNew+n <= new.nkeys())
    if n == 0 {
        return
    }
    // pointers（设置最新的节点指针）
    for i := uint16(0); i < n; i++ { // 更新 idx 列表中的节点指针
        // setPtr(idx, val)
        new.setPtr(dstNew+i, old.getPtr(srcOld+i))
    }
    
    // offsets（更新 offset）
    dstBegins := new.getOffset(dstNew)
    srcBegin := old.getOffset(srcOld)
    for i := uint16(1); i <= n; i++ {
        offset := dstBegin + old.getOffset(srcOld+i) - srcBegin
        new.setOffset(dstNew+i, offset)
    }
    
    // kvs（数据搬移）
    begin := old.kvPos(srcOld)
    end := old.kvPos(srcOld + n)
    copy(new.data[new.kvPos(dstNew):], old.data[begin:end])
}

如图所示，假设以前页中存在 key=1、3、5、7 的四个 KV 对，现在插入 key=6 的 KV 对，那么需要先将 1、3、5 的 KV 先插入到新页中，然后将 key=6 的 KV 对插入到之后，最后将 key=7 的 KV 对插入到末尾，完成一次 newkey 的插入操作。（从这里我们也可以看出为什么推荐使用 auto_increment 的主键索引了，那样就省略了 B+ 树分裂重新构造的开销了；只在页面不够用，才可能触发页的分裂操作）。

NodeAppendKV 函数将一个 KV 对复制到新节点。

// copy a KV into the position
func nodeAppendKV(new BNode, idx uint16, ptr uint64, key []byte, val []byte) {
    // ptrs
    new.setPtr(idx, ptr)
    // KVs
    pos := new.kvPos(idx)
    binary.LittleEndian.PutUint16(new.data[pos+0:], uint16(len(key)))
    binary.LittleEndian.PutUint16(new.data[pos+2:], uint16(len(val)))
    copy(new.data[pos+4:], key)
    copy(new.data[pos+4+uint16(len(key)):], val)
    // the offset of the next key
    new.setOffset(idx+1, new.getOffset(idx) + 4 + uint16((len(key)+len(val))))
}

Step 3: Recursive Insertion

用于插入键的主函数。

// 将 KV 插入节点，结果可能会被分割成两个节点。调用者负责取消输入节点的分配，并分割和分配结果节点。
func treeInsert(tree *BTree, node BNode, key []byte, val []byte) BNode {
    // 结果节点。它允许大于1页，如果大于1页，将被拆分
    new := BNode{data: make([]byte, 2 * BTREE_PAGE_SIZE)}
    
    // where to insert the key?
    idx := nodeLookupLE(node, key)
    // 根据节点类型进行操作
    switch node.btype() {
        case BNODE_LEAF: // 可能是更新或者插入操作
            if bytes.Equal(key, node.getKey(idx)) {
                // found the key, update it
                leafUpdate(new, node, key, val)
            } else {
                // insert it after the position
                leafInsert(new, node, idx+1, key, val)
            }
       case BNODE_NODE:
           // internal node, insert it to a kid node.
           nodeInsert(tree, new, node, idx, key, val)
       default:
           panic("bad node")
    }
    return new
}

leafUpdate 函数与 leafInsert 函数类似。

Step 4: Handle Internal Nodes

现在是处理内部节点的代码。

// KV insertion to an internal node
func nodeInsert(tree *BTree, new BNode, node BNode, idx uint16, key []byte, val []byte) {
    // get and deallocate the kid node
    kptr := node.getPtr(idx)
    knode := tree.get(kptr)
    tree.del(kptr)
    // 递归插入进 kid node
    knode = treeInsert(tree, knode, key, val)
    // split the result
    nsplit, splited := nodeSplit3(knode)
    // update the kid links
    nodeReplaceKidN(tree, new, node, idx, splited[:nsplit]...)
}
 

Step 5: Split Big Nodes

将键插入节点会增加其大小，导致其超出页面大小。在这种情况下，节点被分割成多个更小的节点。

允许的最大键大小和值大小仅保证单个 KV 对始终适合一页。在最坏的情况下，胖节点被分成3个节点（中间有一个大的 KV 对）。

// split a bigger-than-allowed node into two.
// the second node always fits on a page.
func nodeSplit2(left BNode, right BNode, old BNode) {
    // code omitted...
}
 
// split a node if it's too big. the results are 1~3 nodes.
func nodeSplit3(old BNode) (uint16, [3]BNode) {
    if old.nbytes() <= BTREE_PAGE_SIZE {
        old.data = old.data[:BTREE_PAGE_SIZE]
        return 1, [3]BNode{old}
    }
    left := BNode{make([]byte, 2*BTREE_PAGE_SIZE)} // might be split later
    right := BNode{make([]byte, BTREE_PAGE_SIZE)}
    nodeSplit2(left, right, old)
    if left.nbytes() <= BTREE_PAGE_SIZE {
        left.data = left.data[:BTREE_PAGE_SIZE]
        return 2, [3]BNode{left, right}
    }
    // the left node is still too large
    leftleft := BNode{make([]byte, BTREE_PAGE_SIZE)}
    middle := BNode{make([]byte, BTREE_PAGE_SIZE)}
    nodeSplit2(leftleft, middle, left)
    assert(leftleft.nbytes() <= BTREE_PAGE_SIZE)
    return 3, [3]BNode{leftleft, middle, right}
}

Step 6: Update Internal Nodes

将键插入节点可能会产生 1、2 或 3 个节点。父节点必须相应地更新自身。更新内部节点的代码与更新叶节点的代码类似。

// replace a link with multiple links
func nodeReplaceKidN( 
    tree *BTree, new BNode, old BNode, idx uint16,
    kids ...BNode,
 ) {
    inc := uint16(len(kids))
    new.setHeader(BNODE_NODE, old.nkeys()+inc-1)
    nodeAppendRange(new, old, 0, 0, idx)
    for i, node := range kids {
        nodeAppendKV(new, idx+uint16(i), tree.new(node), node.getKey(0), nil)
    }
    nodeAppendRange(new, old, idx+inc, idx+1, old.nkeys()-(idx+1))
}

我们已经完成了 B 树的插入，删除以及其余代码将在下一章介绍。