B Tree Index を実装する - それが僕には楽しかったんです。

この記事は「けんつの１人 DBMS アドベントカレンダー Advent Calendar 2019 - Adventar」 15 日目の記事です。

はじめに
B Tree を実装する
インデックスの色々
- クラスタインデックス
- セカンダリインデックス
B Tree アルゴリズム
終わりに

はじめに

どうも、最近眼精疲労がひどすぎるけど眼球ってどうやって休めたらいいかわからないけんつです。

今日明日は現状最も苦戦した B Tree Index と Buffer Pool の実装話になります。

B Tree を実装する

今回は B+ Tree でなく B Tree を実装する。また削除は考慮しない。
完全に考慮しない訳ではなくて一応削除を実装することを考慮はしている。ただ物理削除ではなく論理削除としている。
というのもなぜかと言うと、B Tree のなかで最もコストの高い動作が発生するのが物理削除であるという理由から。

どのように論理削除するというとフラグを立てて探索をスキップするという単純なものを実装している。

インデックスの色々

インデックスと言っても色々あってここではその話をする。あくまでも MySQL が前提。

クラスタインデックス

MySQL のインデックスには２種類あってクラスタインデックスとセカンダリインデックスというものが存在する。クラスタインデックスは雑に言うと主キーで構成されたインデックスで暗黙的に構築される。 PK が定義されていないなら全てのカラムから NOT NULL な UNIQUE インデックスを最初に検索し、それを元に構成する。それすらもない場合は InnoDB が内部的に行に自動で付与されるインクリメントカラムを追加しその値で構成する。

セカンダリインデックス

クラスタインデックス以外のインデックスを全てセカンダリインデックスと呼ぶ。これは特定のカラムに貼るものと思って間違いは無いはず。
そして実際にセカンダリインデックスとクラスタインデックスがどのように関連しているかというと、セカンダリインデックスで探索するものは主キー、つまりクラスタインデックスであることが多く
セカンダリインデックスで探すべき主キーを見つけ、そのあとクラスタインデックスを探索してその主キーに紐付いている行を見つけている。
しかし、セカンダリインデックスに必要な情報が全て格納されていれば、クラスタインデックスによる探索をスキップできるため高速に動作させることができる。

B Tree アルゴリズム

「B Tree 実装」とかってググるとなぜか MySQL のインデックスなどが引っかかってきてハイパーわかりにくかったのでアルゴリズムの手順的なものは以下のスライドを参考にした。

speakerdeck.com

これを見れば大体行ける。
そして、以下のデモサイトを見ると困った時に B Tree がどのように処理を行っていくかが試せるからすごく良かった。
www.cs.usfca.edu

実装方針

B+ Tree ではなく B Tree
物理削除は実装しない
数値と文字列*1に対応させる
クラスタインデックスとしての利用を想定

実装するにあたっての知見・感想

アイテムの閾値

ノードが最大何個までアイテムを持つかという値を決める必要があるが、これは間違いなく奇数がいい。
なぜかというと、閾値を 3 にした場合を例に説明するならばノードが持つアイテムが３つになった段階で最大値・最小値を2番目のアイテムの子ノードにする必要があるから。
偶数にする場合はまた何かしら別に処理する必要があるっぽいが単純に面倒くさいので 3, 5, 7 あたりの真ん中の値が一意に求まるものにすると実装しやすい。

ラッチの実装

CMU の講義でも言われているが、ノードを探索、追加、更新する場合にはラッチを利用するのがよいと言われている。
例えば探索なら、まずルートノードのロック*2をとり、探索する前に子ノードのロックも確保する。
ルートノードの探索が終わったら、ロックを開放して子ノードの更に子ノードのロックを取得する、というプロセスを値が見つかるまで繰り返す。

今回は Github の PR でもわかる通り、今回はラッチをノードレベルではなくインデックスのツリー構造単位でかけている。
これはめちゃくちゃパフォーマンスが悪いのは知っているけど、↑のようにロックをノードレベルで取得がめちゃくちゃめんどくさくて*3 そうなってしまった。

Golang の interface を上手く使う

ノードが保持するアイテムを Item という以下の interface として宣言している

type Item interface {
        Less(itm Item) bool
        Equal(itm Item) bool
        IsSkip() bool
}

そのため B Tree で使うデータ型は上記のメソッドを実装していればいいから、データ型によって異なる大小比較とかが抽象化できて
B Tree 自体のコードがあまり複雑にならなかった。
データ型の宣言は以下の様になっている。

package storage


type IntItem struct {
        Value int32
        Delete bool
}
type StringItem struct {
        Value string
        Delete bool
}

type Item interface {
        Less(itm Item) bool
        Equal(itm Item) bool
        IsSkip() bool
}

func (i IntItem) Less(itm Item) bool {
        v, ok := itm.(IntItem)
        if !ok {
                return false
        }

        return i.Value < v.Value
}

func (i IntItem) Equal(itm Item) bool {
        v, ok := itm.(IntItem)
        if !ok {
                return false
        }
        return i.Value == v.Value
}

func (i IntItem) IsSkip() bool {
        return i.Delete
}

func (s StringItem) Less(itm Item) bool {
        v, ok := itm.(StringItem)
        if !ok {
                return false
        }
        return s.Value < v.Value
}

func (s StringItem) Equal(itm Item) bool {
        v, ok := itm.(StringItem)
        if !ok {
                return false
        }
        return s.Value == v.Value
}

func (s StringItem) IsSkip() bool {
        return s.Delete
}

B Tree の実装

もう解説がめんどくさくなったので、これは本当に自信が無いが*4 コード載せて終わりとする

package storage

import (
        "sort"
        "sync"
)

const (
        MergeThreshold = 3
)

type Items []Item

type Node struct {
        Items Items
        Children []*Node
}

type BTree struct {
        Top *Node
        Mutex sync.RWMutex
}

func NewBTree() *BTree {
        return &BTree{
                Top: nil,
                Mutex: sync.RWMutex{},
        }
}

func (b *BTree) Insert(item Item) {

        b.Mutex.Lock()
        defer b.Mutex.Unlock()

        if b.Top == nil {
                node := new(Node)
                node.Items.InsertAt(item, 0)
                b.Top = node
        } else {
                b.Top.Insert(item)
        }
}

func (n *Node) Insert(item Item) {

        _, index := n.Items.Find(item)

        if len(n.Children) == 0 {

                n.Items.InsertAt(item, index)

                if len(n.Items) == MergeThreshold {
                        n.Split()
                }
                return
        }

        if len(n.Children[index].Items) == MergeThreshold - 1 {
                n.SplitChildren(item, index)

                if len(n.Items) == MergeThreshold {
                        n.Split()
                }
                return
        }

        n.Children[index].Insert(item)
}

func (n *Node) SplitChildren(item Item, index int) {

        _, i := n.Children[index].Items.Find(item)
        n.Children[index].Items.InsertAt(item, i)

        leftItem := n.Children[index].Items[MergeThreshold/2-1]
        middleItem := n.Children[index].Items[MergeThreshold/2]
        rightItem := n.Children[index].Items[MergeThreshold/2+1]

        n.Children = append(n.Children[:index], n.Children[index+1:]...)

        _, middleIndex := n.Items.Find(middleItem)
        n.Items.InsertAt(middleItem, middleIndex)

        leftNode := new(Node)
        rightNode := new(Node)

        leftNode.Insert(leftItem)
        rightNode.Insert(rightItem)

        n.Children = append(n.Children, leftNode)
        n.Children = append(n.Children, rightNode)

        sort.Slice(n.Children, func(i, j int) bool {
                return n.Children[i].Items[0].Less(n.Children[j].Items[0])
        })

}

func (n *Node) Split() {

        left := new(Node)
        middle := n.Items[MergeThreshold/2]
        right := new(Node)

        left.Items.InsertAt(n.Items[MergeThreshold/2-1], 0)
        right.Items.InsertAt(n.Items[MergeThreshold/2+1], 0)
        n.Items = append([]Item{}, middle)

        if len(n.Children) == MergeThreshold + 1 {

                var nodes []*Node

                left.Children = append(left.Children, []*Node{n.Children[0], n.Children[1]}...)
                right.Children = append(right.Children, []*Node{n.Children[2], n.Children[3]}...)

                nodes = append(nodes, left)
                nodes = append(nodes, right)
                n.Children = nodes
                return
        }

        n.Children = append(n.Children, left)
        n.Children = append(n.Children, right)
}

func (b *BTree) Find(item Item) (bool, int) {

        b.Mutex.RLock()
        defer b.Mutex.RUnlock()

        if b.Top == nil {
                return false, -1
        }

        return b.Top.Find(item)
}

func (n *Node) Find(item Item) (bool, int) {
        ok, index := n.Items.Find(item)
        if ok {
                return ok, index
        }

        if len(n.Children) == 0 {
                return false, -1
        }

        return n.Children[index].Find(item)
}

func (i *Items) Find(item Item) (bool, int) {
        for index,itm := range *i {

                if itm.IsSkip() {
                        continue
                }

                if itm.Equal(item) {
                        return true, index
                }

                if !itm.Less(item) {
                        return false, index
                }

        }
        return false, len(*i)
}

func (i *Items) InsertAt(item Item, index int) {
        *i = append(*i, nil)
        if index < len(*i) {
                copy((*i)[index+1:], (*i)[index:])
        }
        (*i)[index] = item
}

func (b *BTree) Get(item Item) Item {

        b.Mutex.RLock()
        defer b.Mutex.RUnlock()

        if b.Top == nil {
                return nil
        }

        return b.Top.Get(item)
}

func (n *Node) Get(item Item) Item {
        found, i := n.Items.Find(item)
        if found {
                return n.Items[i]
        } else if len(n.Children) > 0 {
                return n.Children[i].Get(item)
        }
        return nil
}

終わりに

次は、インデックスより苦戦した Buffer Pool について

*1:文字列は辞書順で比較

*2:ここでいうロックとは同時実行制御のロックではなくマルチスレッドの排他処理のロックをイメージして欲しい

*3:INSERT, UPDATE 時のような更新、追加処理のラッチはもっと面倒くさい

*4:テストが本当の正常系しか担保できていない