From c4b7b2e97b9941313eab4dfdf1668c112b8b1320 Mon Sep 17 00:00:00 2001
From: Lucas <junsuny789@gmail.com>
Date: Thu, 12 Sep 2024 14:54:27 -0400
Subject: [PATCH] add radix tree benchmark

---
 CMakeLists.txt          |  2 ++
 external/CMakeLists.txt | 23 +++++++++++++++++---
 script/bench.py         |  2 +-
 src/bliss/bench_rax.h   | 47 +++++++++++++++++++++++++++++++++++++++++
 src/bliss_bench.cpp     |  7 +++++-
 5 files changed, 76 insertions(+), 5 deletions(-)
 create mode 100644 src/bliss/bench_rax.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 1a93347..be7178b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -44,6 +44,7 @@ add_library(bliss OBJECT
     ${CMAKE_SOURCE_DIR}/src/bliss/bliss_index.h
     ${CMAKE_SOURCE_DIR}/src/bliss/bench_lipp.h
     ${CMAKE_SOURCE_DIR}/src/bliss/bench_alex.h
+    ${CMAKE_SOURCE_DIR}/src/bliss/bench_rax.h
     ${CMAKE_SOURCE_DIR}/src/bliss/bench_btree.h
 )
 
@@ -57,6 +58,7 @@ target_link_libraries(bliss PUBLIC
     alex
     lipp
     tlx
+    rax
 )
 
 target_include_directories(bliss PUBLIC
diff --git a/external/CMakeLists.txt b/external/CMakeLists.txt
index 3531b2f..7cfe73d 100644
--- a/external/CMakeLists.txt
+++ b/external/CMakeLists.txt
@@ -31,6 +31,21 @@ FetchContent_MakeAvailable(cxxopts)
 
 
 
+FetchContent_Declare(
+    rax
+    GIT_REPOSITORY https://github.com/antirez/rax
+    GIT_TAG master
+)
+FetchContent_GetProperties(rax)
+if (NOT rax_POPULATED)
+    FetchContent_Populate(rax)
+endif()
+
+add_library(rax INTERFACE)
+target_include_directories(rax INTERFACE ${rax_SOURCE_DIR})
+
+
+
 FetchContent_Declare(
     alex
     GIT_REPOSITORY https://github.com/microsoft/ALEX.git
@@ -59,10 +74,12 @@ endif()
 add_library(lipp INTERFACE)
 target_include_directories(lipp INTERFACE ${lipp_SOURCE_DIR}/src/core)
 
+
+
 FetchContent_Declare(
-        tlx
-        GIT_REPOSITORY https://github.com/tlx/tlx.git
-        GIT_TAG master
+    tlx
+    GIT_REPOSITORY https://github.com/tlx/tlx.git
+    GIT_TAG master
 )
 FetchContent_GetProperties(tlx)
 if (NOT tlx_POPULATED)
diff --git a/script/bench.py b/script/bench.py
index efaac72..b380d78 100755
--- a/script/bench.py
+++ b/script/bench.py
@@ -7,7 +7,7 @@
 from infra.pybliss import BlissArgs, PyBliss
 from infra.util import get_file_params
 
-INDEXES = ["btree"]
+INDEXES = ["btree", "radix_tree"]
 PRELOAD_FACTOR = 0.4
 WRITE_FACTOR = 0.4
 READ_FACTOR = 0.2
diff --git a/src/bliss/bench_rax.h b/src/bliss/bench_rax.h
new file mode 100644
index 0000000..83416b0
--- /dev/null
+++ b/src/bliss/bench_rax.h
@@ -0,0 +1,47 @@
+#ifndef BLISS_BENCH_RAX
+#define BLISS_BENCH_RAX
+
+#include <vector>
+
+#include "bliss/bliss_index.h"
+
+#include "rax.h"
+
+
+namespace bliss {
+
+template <typename KEY_TYPE, typename VALUE_TYPE>
+class BlissRaxIndex : public BlissIndex<KEY_TYPE, VALUE_TYPE> {
+   public:
+    rax* _index;
+    BlissRaxIndex() {
+        _index = raxNew();
+    };
+
+    void bulkload(
+        std::vector<std::pair<KEY_TYPE, VALUE_TYPE>> values) override {
+        // expects the pairs to be pre-sorted before performing bulk load
+        for (const auto& pair : values) {
+            put(pair.first, pair.second);
+        }
+    }
+
+    bool get(KEY_TYPE key) override { 
+        std::string keyStr = std::to_string(key);
+        void* result = raxFind(_index, (unsigned char*)keyStr.c_str(), keyStr.size());
+        return (result == raxNotFound) ? false : true;
+     }
+
+    void put(KEY_TYPE key, VALUE_TYPE value) override {
+        std::string keyStr = std::to_string(key);
+        raxInsert(_index, (unsigned char*)keyStr.c_str(), keyStr.size(), (void*)new VALUE_TYPE(value), NULL);
+    }
+
+    void end_routine() override {
+        raxFree(_index);
+    }
+};
+
+}  // namespace bliss
+
+#endif  // !BLISS_BENCH_RAX
diff --git a/src/bliss_bench.cpp b/src/bliss_bench.cpp
index d00d8a0..d688712 100644
--- a/src/bliss_bench.cpp
+++ b/src/bliss_bench.cpp
@@ -1,5 +1,7 @@
 #include <alex.h>
 #include <lipp.h>
+
+#include "rax.h"
 #include <spdlog/common.h>
 
 #include <cxxopts.hpp>
@@ -9,6 +11,7 @@
 #include "bliss/bench_alex.h"
 #include "bliss/bench_btree.h"
 #include "bliss/bench_lipp.h"
+#include "bliss/bench_rax.h"
 #include "bliss/bliss_index.h"
 #include "bliss/util/reader.h"
 #include "bliss/util/timer.h"
@@ -49,7 +52,7 @@ BlissConfig parse_args(int argc, char *argv[]) {
             cxxopts::value<int>()->default_value("0"))(
             "v,verbosity", "Verbosity [0: Info| 1: Debug | 2: Trace]",
             cxxopts::value<int>()->default_value("0")->implicit_value("1"))(
-            "i,index", "Index type [alex | lipp | btree | bepstree | lsm]",
+            "i,index", "Index type [alex | lipp | btree | bepstree | lsm | radix_tree]",
             cxxopts::value<std::string>()->default_value("btree"))(
             "file_type", "Input file type [binary | txt]",
             cxxopts::value<std::string>()->default_value("txt"))(
@@ -264,6 +267,8 @@ int main(int argc, char *argv[]) {
         index.reset(new bliss::BlissLippIndex<key_type, value_type>());
     } else if (config.index == "btree") {
         index.reset(new bliss::BlissBTreeIndex<key_type, value_type>());
+    } else if (config.index == "radix_tree") {
+        index.reset(new bliss::BlissRaxIndex<key_type, value_type>());
     } else {
         spdlog::error("{} not implemented yet", config.index);
     }