diff --git a/src/seriousdb/b_plus_tree.py b/src/seriousdb/b_plus_tree.py new file mode 100644 index 0000000..3081a01 --- /dev/null +++ b/src/seriousdb/b_plus_tree.py @@ -0,0 +1,277 @@ +"""b+tree implementation for datalookups in O(log n). + +this module will hold the b+tree datastructure for loading and deleting data in and out of memory +!IMPORTANT this implementation is not Threadsafe +""" + +from __future__ import annotations + + +class _Node: + """internal class representing a node from the b+Tree.""" + + _parent: _InternalNode | None + _keys: list[int] + + def __init__(self) -> None: + self._parent = None + self._keys = [] + + +class _LeafNode(_Node): + """internal class representing a _leafnode from the b+Tree.""" + + _values: list[bytes] + _next_leaf: _LeafNode | None + + def __init__(self) -> None: + super().__init__() + self._values = [] + self._next_leaf = None + + +class _InternalNode(_Node): + """internal class representing.""" + + _children: list[_Node] + + def __init__(self) -> None: + super().__init__() + self._children = [] + + +class BPlusTree: + """a Btree like Datastructure to retrieve and put data in O(log n) time. + + the public functions will firstly be get() put() delete() range() + + !IMPORTANT: this structure is currently not threadsafe. + + Parameters + ---------- + _maxChildren :int + the tree takes in the maximum number of children a internal node should have. + + Attributes + ---------- + _maxChildren : int + The maximum number of children (mostly in literature visualized as m) an internal Node can point to + _maxKeys: int + The maximum number of Keys a _Node can hold derrived from the maximum number of children (m -1) + _root : _Node or None + the root Node of the tree the first will be a leafNode and after enough insertions mostly an internal Node + + Raises + ------ + ValueError if the maxchildren parameter is below 3 + """ + + _maxChildren: int = 0 + _maxKeys: int = 0 + _root: _LeafNode | _InternalNode | None + _size: int = 0 + + def __init__(self, maxChildren: int) -> None: + if maxChildren < 3: + raise ValueError("maxChildren must be at least 3") + self._maxChildren = maxChildren + self._maxKeys = self._maxChildren - 1 + self._root = None + self._size = 0 + + def put(self, key: int, value: bytes) -> bool: + """Put an key value pair into the tree. + + if the tree is empty a leafNode is created at root and filled with the first key and value. + if the tree alereay has some nodes the search starts. The next node from the children gets chosen via the comparator keys. + If a leafnode has to many values or keys put splits the leafnote. + + Parameters + ---------- + key : int + the key to access the value. + value : bytes + bytes (Pages) to store. + + Returns + ------- + is_new_key : bool + ``True`` if `key` did not exist before, ``False`` if an existing + value was replaced. + """ + is_new_key = True + self._size += 1 + if self._root is None: + self._root = self._createLeaf(key, value) + return is_new_key + + node: _LeafNode | _InternalNode | _Node = self._root + while True: + if isinstance(node, _LeafNode): + if self._isKeyInLeaf(node, key): + key_index = node._keys.index(key) + node._values[key_index] = value + return False + + # NOTE: this find_position call only works because we check for duplicates in the if case before maby for later its more smart to make 2 helpers out of this + insert_index: int = self._find_position(node._keys, key) + + node._keys.insert(insert_index, key) + node._values.insert(insert_index, value) + + self._check_split(node) + return True + elif isinstance(node, _InternalNode): + internalNode: _InternalNode = node + nodeIndex = self._find_position(internalNode._keys, key) + node = internalNode._children[nodeIndex] + + def get(self, key: int) -> bytes | None: + """Return a value from the given key. if the key does not exist in the tree this function returns None. + + Parameters + ---------- + key : int + the key to access the value. + + Returns + ------- + value : bytes | None + returns bytes if the key is in the tree otherwise this function returns None + """ + value: bytes | None = None + if self._root is None: + return value + + node: _InternalNode | _LeafNode | _Node = self._root + while isinstance(node, _InternalNode): + nodeIndex: int = self._find_position(node._keys, key) + node = node._children[nodeIndex] + + # node should always be a _LeafNode this might be to defensive + if isinstance(node, _LeafNode) and key in node._keys: + keyIndex: int = node._keys.index(key) + value = node._values[keyIndex] + + return value + + def size(self) -> int: + """Return the current keyvaluePairs in the tree. + + Returns + ------- + value : int + returns the number of keyvalue pairs in the tree + """ + return self._size + + def delete(self, key: int) -> bool: + """Delete algorithm deletes key value pair.""" + return True + + """------ INTERNAL FUNCTIONS ------""" + + def _createLeaf(self, key: int, value) -> _LeafNode: + node = _LeafNode() + node._keys.append(key) + key_index = node._keys.index(key) + node._values.insert(key_index, value) + return node + + def _createInternal(self, key, child1, child2) -> _InternalNode: + internalNode: _InternalNode = _InternalNode() + internalNode._keys.append(key) + internalNode._children.append(child1) + internalNode._children.append(child2) + + return internalNode + + def _isKeyInLeaf(self, node: _LeafNode, key: int) -> bool: + return key in node._keys + + def _find_position(self, keys: list[int], key: int) -> int: + insert_index: int = 0 + + while insert_index < len(keys) and keys[insert_index] <= key: + insert_index += 1 + + return insert_index + + def _split_leaf(self, node: _LeafNode) -> None: + newLeaf: _LeafNode = _LeafNode() + keyLen: int = len(node._keys) + splitIndex: int = keyLen // 2 + + newLeaf._keys = node._keys[splitIndex:] + newLeaf._values = node._values[splitIndex:] + + node._keys = node._keys[:splitIndex] + node._values = node._values[:splitIndex] + + newLeaf._next_leaf = node._next_leaf + node._next_leaf = newLeaf + + newSeperator: int = newLeaf._keys[0] + + if node._parent is None: + # if the parent of an leafnote is None it is guaranteed that the leafnote lies at root. + internalNode: _InternalNode = self._createInternal( + newSeperator, node, newLeaf + ) + + newLeaf._parent = internalNode + node._parent = internalNode + self._root = internalNode + else: + parent: _InternalNode = node._parent + + child_index: int = parent._children.index(node) + parent._children.insert(child_index + 1, newLeaf) + parent._keys.insert(child_index, newSeperator) + newLeaf._parent = parent + + self._check_split(parent) + + def _split_internal(self, node: _InternalNode) -> None: + oldNode = node + newNode = _InternalNode() + + keyLen: int = len(node._keys) + splitIndex: int = keyLen // 2 + + tempKey = oldNode._keys[splitIndex] + newNode._keys = oldNode._keys[splitIndex + 1 :] + oldNode._keys = oldNode._keys[:splitIndex] + + newNode._children = oldNode._children[splitIndex + 1 :] + oldNode._children = oldNode._children[: splitIndex + 1] + + for child in newNode._children: + child._parent = newNode + + if oldNode._parent is None: + internalNode: _InternalNode = self._createInternal( + tempKey, oldNode, newNode + ) + + newNode._parent = internalNode + oldNode._parent = internalNode + self._root = internalNode + else: + parent: _InternalNode = oldNode._parent + + child_index: int = parent._children.index(oldNode) + parent._children.insert(child_index + 1, newNode) + parent._keys.insert(child_index, tempKey) + newNode._parent = parent + + self._check_split(parent) + + def _check_split(self, node: _Node | _InternalNode | _LeafNode) -> None: + if len(node._keys) <= self._maxKeys: + return + + if isinstance(node, _LeafNode): + self._split_leaf(node) + elif isinstance(node, _InternalNode): + self._split_internal(node) diff --git a/tests/test_b_plus_tree.py b/tests/test_b_plus_tree.py new file mode 100644 index 0000000..38fdae9 --- /dev/null +++ b/tests/test_b_plus_tree.py @@ -0,0 +1,95 @@ +"""tests for the b+Tree implementation""" + +from seriousdb.b_plus_tree import BPlusTree, _InternalNode, _LeafNode, _Node + + +def assertTreeStructure(tree: BPlusTree) -> None: + root = tree._root + + if root is None: + return + # okay here i will just make it a really long for loop so it cannot run infenetly if there is an error in the tree code + stack = [root] + while len(stack) > 0: + node: _Node = stack.pop() + + assert len(node._keys) <= tree._maxKeys + + if node is not tree._root: + assert node._parent is not None + + if isinstance(node, _LeafNode): + assert len(node._keys) == len(node._values) + + if isinstance(node, _InternalNode): + assert len(node._children) == len(node._keys) + 1 + stack.extend(node._children) + + +def test_put_first_key(): + maxChildren: int = 3 + tree: BPlusTree = BPlusTree(maxChildren) + + result = tree.put(1, b"2") + + assert result is True + + +def test_put_same_key(): + maxChildren: int = 3 + tree: BPlusTree = BPlusTree(maxChildren) + + tree.put(1, b"2") + result: bool = tree.put(1, b"2") + + assert result is False + + +def test_force_split(): + maxChildren: int = 3 + tree: BPlusTree = BPlusTree(maxChildren) + tree.put(1, b"AA") + tree.put(2, b"AA") + tree.put(3, b"AA") + tree.put(4, b"AA") + + assertTreeStructure(tree) + assert isinstance(tree._root, _InternalNode) + + +def test_stress_put(): + maxChildren: int = 3 + tree: BPlusTree = BPlusTree(maxChildren) + for i in range(0, 100000, 1): + tree.put(i, b"bbb") + + assertTreeStructure(tree) + + +def test_get_success(): + maxChildren: int = 3 + tree: BPlusTree = BPlusTree(maxChildren) + + tree.put(1, b"AA") + value = tree.get(1) + + assert value == b"AA" + + +def test_get_empyTree(): + maxChildren: int = 3 + tree: BPlusTree = BPlusTree(maxChildren) + # retrive + value = tree.get(1) + assert value == None + + +def test_get_KeyNotInTree(): + maxChildren: int = 3 + tree: BPlusTree = BPlusTree(maxChildren) + + tree.put(2, b"AA") + tree.put(3, b"BB") + + value = tree.get(5) + assert value == None