cleanup: benchmark

benchmark/compare.py

@@ -5,31 +5,28 @@
 class DiffSpan:
     def __init__(
         self,
-        spanID: str,
-        operationName: str,
-        duration_head: int,
-        duration_branch: int,
-        tags: list[dict[str, str]],
+        span_id: str,
+        operation_name: str,
+        main_duration: int,
+        branch_duration: int,
     ):
-        self.spanID = spanID
-        self.operationName = operationName
-        self.duration_head = duration_head
-        self.duration_branch = duration_branch
-        self.improvement = f"{round((1 - duration_branch / duration_head) * 100)}%"
-        self.tags = tags
+        self.span_id = span_id
+        self.operation_name = operation_name
+        self.main_span_duration = main_duration
+        self.branch_span_duration = branch_duration
+        self.improvement = f"{round((1 - main_duration / branch_duration) * 100)}%"
         self.children: dict[str, "DiffSpan"] = {}
 
     def add_child(self, span_id: str, child: "DiffSpan"):
         self.children[span_id] = child
 
     def to_dict(self, include_children: bool = True):
         res = {
-            "spanID": self.spanID,
-            "operationName": self.operationName,
-            "duration_head": self.duration_head,
-            "duration_branch": self.duration_branch,
+            "span_id": self.span_id,
+            "operation_name": self.operation_name,
+            "main_duration": self.main_span_duration,
+            "branch_duration": self.branch_span_duration,
             "improvement": self.improvement,
-            "tags": self.tags,
         }
 
         if include_children:
@@ -39,21 +36,22 @@ def to_dict(self, include_children: bool = True):
 
 
 class DiffTree:
-    def __init__(self, head_trace: dict, branch_trace: dict):
-        self.root_span = self.__build_diff_tree(head_trace, branch_trace)
+    def __init__(self, main_trace: dict, branch_trace: dict):
+        self.root_span = self.__build_diff_tree(main_trace, branch_trace)
+
+    def __build_diff_tree(self, main_trace: dict, branch_trace: dict):
+        assert main_trace["operationName"] == branch_trace["operationName"]
 
-    def __build_diff_tree(self, head_trace: dict, branch_trace: dict):
         diff_span = DiffSpan(
-            head_trace["spanID"],
-            head_trace["operationName"],
-            head_trace["duration"],
+            main_trace["spanID"],
+            main_trace["operationName"],
+            main_trace["duration"],
             branch_trace["duration"],
-            head_trace["tags"],
         )
 
         # Recursively build the tree for child spans
         for head_child_span, branch_child_span in zip(
-            head_trace["children"], branch_trace["children"]
+            main_trace["children"], branch_trace["children"]
         ):
             child_span = self.__build_diff_tree(head_child_span, branch_child_span)
             diff_span.add_child(head_child_span["spanID"], child_span)
@@ -73,10 +71,10 @@ def main():
     current_dir = pathlib.Path(__file__).parent.absolute()
 
     for operation in ["pyg_to_arangodb", "arangodb_to_pyg"]:
-        head_trace = json.load(open(f"{current_dir}/traces/head/{operation}.json"))
+        main_trace = json.load(open(f"{current_dir}/traces/main/{operation}.json"))
         branch_trace = json.load(open(f"{current_dir}/traces/branch/{operation}.json"))
 
-        diff_tree = DiffTree(head_trace, branch_trace)
+        diff_tree = DiffTree(main_trace, branch_trace)
         diff_tree.to_json_file(operation)
 
         print("-" * 50)