feat: pretty print B+ tree (#503)

* pretty print B+ tree

* add back comment of `Draw()`

* fix lint

* fix clang-tidy

* adapt to new interface

Signed-off-by: Alex Chi <[email protected]>

* fix

Signed-off-by: Alex Chi <[email protected]>

* fix

Signed-off-by: Alex Chi <[email protected]>

---------

Signed-off-by: Alex Chi <[email protected]>
Co-authored-by: Alex Chi <[email protected]>

Author: kikimo

src/include/storage/index/b_plus_tree.h

@@ -1,13 +1,3 @@
-//===----------------------------------------------------------------------===//
-//
-//                         CMU-DB Project (15-445/645)
-//                         ***DO NO SHARE PUBLICLY***
-//
-// Identification: src/include/index/b_plus_tree.h
-//
-// Copyright (c) 2018, Carnegie Mellon University Database Group
-//
-//===----------------------------------------------------------------------===//
 /**
  * b_plus_tree.h
  *
@@ -22,9 +12,11 @@
 
 #include <algorithm>
 #include <deque>
+#include <iostream>
 #include <optional>
 #include <queue>
 #include <shared_mutex>
+#include <string>
 #include <vector>
 
 #include "common/config.h"
@@ -38,6 +30,8 @@
 
 namespace bustub {
 
+struct PrintableBPlusTree;
+
 /**
  * @brief Definition of the Context class.
  *
@@ -63,6 +57,7 @@ class Context {
 };
 
 #define BPLUSTREE_TYPE BPlusTree<KeyType, ValueType, KeyComparator>
+
 // Main class providing the API for the Interactive B+ Tree.
 INDEX_TEMPLATE_ARGUMENTS
 class BPlusTree {
@@ -102,6 +97,19 @@ class BPlusTree {
   // Draw the B+ tree
   void Draw(BufferPoolManager *bpm, const std::string &outf);
 
+  /**
+   * @brief draw a B+ tree, below is a printed
+   * B+ tree(3 max leaf, 4 max internal) after inserting key:
+   *  {1, 5, 9, 13, 17, 21, 25, 29, 33, 37, 18, 19, 20}
+   *
+   *                               (25)
+   *                 (9,17,19)                          (33)
+   *  (1,5)    (9,13)    (17,18)    (19,20,21)    (25,29)    (33,37)
+   *
+   * @return std::string
+   */
+  auto DrawBPlusTree() -> std::string;
+
   // read data from file and insert one by one
   void InsertFromFile(const std::string &file_name, Transaction *txn = nullptr);
 
@@ -114,11 +122,13 @@ class BPlusTree {
 
   void PrintTree(page_id_t page_id, const BPlusTreePage *page);
 
-  void Dump() {
-    for (auto &f : log) {
-      std::cout << f << std::endl;
-    }
-  }
+  /**
+   * @brief Convert A B+ tree into a Printable B+ tree
+   *
+   * @param root_id
+   * @return PrintableNode
+   */
+  auto ToPrintableBPlusTree(page_id_t root_id) -> PrintableBPlusTree;
 
   // member variable
   std::string index_name_;
@@ -130,4 +140,40 @@ class BPlusTree {
   page_id_t header_page_id_;
 };
 
+/**
+ * @brief for test only. PrintableBPlusTree is a printalbe B+ tree.
+ * We first convert B+ tree into a printable B+ tree and the print it.
+ */
+struct PrintableBPlusTree {
+  int size_;
+  std::string keys_;
+  std::vector<PrintableBPlusTree> children_;
+
+  /**
+   * @brief BFS traverse a printable B+ tree and print it into
+   * into out_buf
+   *
+   * @param out_buf
+   */
+  void Print(std::ostream &out_buf) {
+    std::vector<PrintableBPlusTree *> que = {this};
+    while (!que.empty()) {
+      std::vector<PrintableBPlusTree *> new_que;
+
+      for (auto &t : que) {
+        int padding = (t->size_ - t->keys_.size()) / 2;
+        out_buf << std::string(padding, ' ');
+        out_buf << t->keys_;
+        out_buf << std::string(padding, ' ');
+
+        for (auto &c : t->children_) {
+          new_que.push_back(&c);
+        }
+      }
+      out_buf << "\n";
+      que = new_que;
+    }
+  }
+};
+
 }  // namespace bustub

src/include/storage/page/b_plus_tree_internal_page.h

@@ -11,6 +11,7 @@
 #pragma once
 
 #include <queue>
+#include <string>
 
 #include "storage/page/b_plus_tree_page.h"
 
@@ -72,6 +73,32 @@ class BPlusTreeInternalPage : public BPlusTreePage {
    */
   auto ValueAt(int index) const -> ValueType;
 
+  /**
+   * @brief For test only, return a string representing all keys in
+   * this internal page, formatted as "(key1,key2,key3,...)"
+   *
+   * @return std::string
+   */
+  auto ToString() const -> std::string {
+    std::string kstr = "(";
+    bool first = true;
+
+    // first key of internal page is always invalid
+    for (int i = 1; i < GetSize(); i++) {
+      KeyType key = KeyAt(i);
+      if (first) {
+        first = false;
+      } else {
+        kstr.append(",");
+      }
+
+      kstr.append(std::to_string(key.ToString()));
+    }
+    kstr.append(")");
+
+    return kstr;
+  }
+
  private:
   // Flexible array member for page data.
   MappingType array_[0];

src/include/storage/page/b_plus_tree_leaf_page.h

@@ -10,6 +10,7 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -58,6 +59,31 @@ class BPlusTreeLeafPage : public BPlusTreePage {
   void SetNextPageId(page_id_t next_page_id);
   auto KeyAt(int index) const -> KeyType;
 
+  /**
+   * @brief for test only return a string representing all keys in
+   * this leaf page formatted as "(key1,key2,key3,...)"
+   *
+   * @return std::string
+   */
+  auto ToString() const -> std::string {
+    std::string kstr = "(";
+    bool first = true;
+
+    for (int i = 0; i < GetSize(); i++) {
+      KeyType key = KeyAt(i);
+      if (first) {
+        first = false;
+      } else {
+        kstr.append(",");
+      }
+
+      kstr.append(std::to_string(key.ToString()));
+    }
+    kstr.append(")");
+
+    return kstr;
+  }
+
  private:
   page_id_t next_page_id_;
   // Flexible array member for page data.

src/storage/index/b_plus_tree.cpp

@@ -1,3 +1,4 @@
+#include <sstream>
 #include <string>
 
 #include "common/exception.h"
@@ -292,6 +293,47 @@ void BPLUSTREE_TYPE::ToGraph(page_id_t page_id, const BPlusTreePage *page, std::
   }
 }
 
+INDEX_TEMPLATE_ARGUMENTS
+auto BPLUSTREE_TYPE::DrawBPlusTree() -> std::string {
+  if (IsEmpty()) {
+    return "()";
+  }
+
+  PrintableBPlusTree p_root = ToPrintableBPlusTree(GetRootPageId());
+  std::ostringstream out_buf;
+  p_root.Print(out_buf);
+
+  return out_buf.str();
+}
+
+INDEX_TEMPLATE_ARGUMENTS
+auto BPLUSTREE_TYPE::ToPrintableBPlusTree(page_id_t root_id) -> PrintableBPlusTree {
+  auto root_page_guard = bpm_->FetchPageBasic(root_id);
+  auto root_page = root_page_guard.template As<BPlusTreePage>();
+  PrintableBPlusTree proot;
+
+  if (root_page->IsLeafPage()) {
+    auto leaf_page = root_page_guard.template As<LeafPage>();
+    proot.keys_ = leaf_page->ToString();
+    proot.size_ = proot.keys_.size() + 4;  // 4 more spaces for indent
+
+    return proot;
+  }
+
+  // draw internal page
+  auto internal_page = root_page_guard.template As<InternalPage>();
+  proot.keys_ = internal_page->ToString();
+  proot.size_ = 0;
+  for (int i = 0; i < internal_page->GetSize(); i++) {
+    page_id_t child_id = internal_page->ValueAt(i);
+    PrintableBPlusTree child_node = ToPrintableBPlusTree(child_id);
+    proot.size_ += child_node.size_;
+    proot.children_.push_back(child_node);
+  }
+
+  return proot;
+}
+
 template class BPlusTree<GenericKey<4>, RID, GenericComparator<4>>;
 
 template class BPlusTree<GenericKey<8>, RID, GenericComparator<8>>;

tools/b_plus_tree_printer/b_plus_tree_printer.cpp

@@ -77,7 +77,8 @@ auto main(int argc, char **argv) -> int {
   page_id_t page_id;
   auto header_page = bpm->NewPage(&page_id);
   // create b+ tree
-  BPlusTree<GenericKey<8>, RID, GenericComparator<8>> tree("foo_pk", page_id, bpm, comparator, leaf_max_size, internal_max_size);
+  BPlusTree<GenericKey<8>, RID, GenericComparator<8>> tree("foo_pk", page_id, bpm, comparator, leaf_max_size,
+                                                           internal_max_size);
   // create transaction
   auto *transaction = new Transaction(0);
   while (!quit) {
@@ -110,6 +111,7 @@ auto main(int argc, char **argv) -> int {
         quit = true;
         break;
       case 'p':
+        std::cout << tree.DrawBPlusTree() << std::endl;
         tree.Print(bpm);
         break;
       case 'g':