diff --git a/src/main/kotlin/kestrel/store/btree/BPlusTree.kt b/src/main/kotlin/kestrel/store/btree/BPlusTree.kt new file mode 100644 index 0000000..3c9083d --- /dev/null +++ b/src/main/kotlin/kestrel/store/btree/BPlusTree.kt @@ -0,0 +1,153 @@ +package kestrel.store.btree + +import kestrel.KestrelException +import kestrel.store.codec.KeyCodec +import kestrel.store.page.Page +import kestrel.store.page.PageStore + +class BTreeCorruptionException(message: String) : KestrelException(message) + +/** + * A disk-backed B+Tree mapping order-preserving byte keys (see [KeyCodec]) to + * row pointers. Nodes live on [Page]s fetched through a [BufferPool]; leaves + * are linked for ordered range scans. Insert splits an overflowing node and + * propagates a promoted separator key upward, growing the root when needed. + * Delete removes the entry and, when a node drops below + * [UNDERFLOW_THRESHOLD], redistributes an entry from a sibling with spare + * capacity or merges with one, shrinking the root when it is left with a + * single child. + */ +class BPlusTree private constructor(private val bufferPool: PageStore, rootPageId: Long) { + companion object { + private const val UNDERFLOW_THRESHOLD = Page.PAGE_SIZE / 4 + + fun create(bufferPool: PageStore): BPlusTree { + val page = bufferPool.newPage() + NodeCodec.writeTo(page, BTreeNode.Leaf(page.id, mutableListOf(), mutableListOf(), NodeCodec.NO_PAGE)) + bufferPool.unpinPage(page.id, true) + return BPlusTree(bufferPool, page.id) + } + + fun open(bufferPool: PageStore, rootPageId: Long): BPlusTree = BPlusTree(bufferPool, rootPageId) + } + + var rootPageId: Long = rootPageId + private set + + fun search(key: ByteArray): Long? { + var pageId = rootPageId + while (true) { + val (_, node) = fetch(pageId) + when (node) { + is BTreeNode.Internal -> { + val nextId = node.children[childIndexFor(node.keys, key)] + bufferPool.unpinPage(pageId, false) + pageId = nextId + } + is BTreeNode.Leaf -> { + val idx = lowerBound(node.keys, key) + val result = if (idx < node.keys.size && KeyCodec.compare(node.keys[idx], key) == 0) { + node.values[idx] + } else { + null + } + bufferPool.unpinPage(pageId, false) + return result + } + } + } + } + + fun insert(key: ByteArray, value: Long) { + val split = insertRecursive(rootPageId, key, value) + if (split != null) { + val newRoot = bufferPool.newPage() + NodeCodec.writeTo( + newRoot, + BTreeNode.Internal(newRoot.id, mutableListOf(split.promotedKey), mutableListOf(rootPageId, split.newPageId)) + ) + bufferPool.unpinPage(newRoot.id, true) + rootPageId = newRoot.id + } + } + + fun delete(key: ByteArray): Boolean { + val found = deleteRecursive(rootPageId, key) + val (_, rootNode) = fetch(rootPageId) + if (rootNode is BTreeNode.Internal && rootNode.keys.isEmpty()) { + val onlyChild = rootNode.children[0] + bufferPool.unpinPage(rootPageId, false) + bufferPool.deletePage(rootPageId) + rootPageId = onlyChild + } else { + bufferPool.unpinPage(rootPageId, false) + } + return found + } + + /** Every page belonging to this tree: every internal node and every leaf, reached by a full traversal from the root. */ + fun allPageIds(): List { + val ids = mutableListOf() + val pending = ArrayDeque() + pending.add(rootPageId) + while (pending.isNotEmpty()) { + val pageId = pending.removeFirst() + ids.add(pageId) + val (_, node) = fetch(pageId) + bufferPool.unpinPage(pageId, false) + if (node is BTreeNode.Internal) pending.addAll(node.children) + } + return ids + } + + /** Ordered scan over `[from, to]` (or `[from, to)` when [toInclusive] is false). Null bounds are open. */ + fun rangeScan(from: ByteArray? = null, to: ByteArray? = null, toInclusive: Boolean = false): Iterator> { + val startLeafId = leftmostLeafFor(from) + return iterator { + var currentLeafId = startLeafId + var firstLeaf = true + outer@ while (currentLeafId != NodeCodec.NO_PAGE) { + val (_, leaf) = fetch(currentLeafId) + leaf as BTreeNode.Leaf + var index = if (firstLeaf && from != null) lowerBound(leaf.keys, from) else 0 + firstLeaf = false + while (index < leaf.keys.size) { + val currentKey = leaf.keys[index] + if (to != null) { + val cmp = KeyCodec.compare(currentKey, to) + if (cmp > 0 || (cmp == 0 && !toInclusive)) { + bufferPool.unpinPage(currentLeafId, false) + break@outer + } + } + yield(currentKey to leaf.values[index]) + index++ + } + val next = leaf.nextLeaf + bufferPool.unpinPage(currentLeafId, false) + currentLeafId = next + } + } + } + + private fun leftmostLeafFor(from: ByteArray?): Long { + var pageId = rootPageId + while (true) { + val (_, node) = fetch(pageId) + when (node) { + is BTreeNode.Internal -> { + val idx = if (from == null) 0 else childIndexFor(node.keys, from) + val nextId = node.children[idx] + bufferPool.unpinPage(pageId, false) + pageId = nextId + } + is BTreeNode.Leaf -> { + bufferPool.unpinPage(pageId, false) + return pageId + } + } + } + } + + private data class SplitResult(val promotedKey: ByteArray, val newPageId: Long) +