cleanup compare.py

benchmark/compare.py

@@ -1,85 +1,137 @@
 import json
 import pathlib
+from typing import Optional
+
+
+def sort_children_by_start_time(children):
+    return sorted(children, key=lambda span: span["startTime"])
+
+
+def compare_span(master_child: Optional[dict], branch_child: Optional[dict]):
+    if master_child and branch_child:
+        assert master_child.get("operationName") == branch_child.get("operationName")
+        assert master_child.get("tags") == branch_child.get("tags")
+
+    operation_name = (
+        master_child.get("operationName")
+        if master_child
+        else branch_child.get("operationName")
+    )
+
+    master_duration = master_child.get("duration") if master_child else None
+    branch_duration = branch_child.get("duration") if branch_child else None
+    improvement = (
+        f"{round((1 - branch_duration / master_duration) * 100)}%"
+        if master_duration and branch_duration
+        else None
+    )
+
+    comparison = {
+        "operationName": operation_name,
+        "master_duration": master_duration,
+        "branch_duration": branch_duration,
+        "improvement": improvement,
+        "tags": master_child.get("tags") if master_child else branch_child.get("tags"),
+        "children": [],
+    }
+
+    if master_child and branch_child:
+        comparison["children"] = compare_children(
+            master_child["children"], branch_child["children"]
+        )
 
+    return comparison
 
-class DiffSpan:
-    def __init__(
-        self,
-        span_id: str,
-        operation_name: str,
-        master_duration: int,
-        branch_duration: int,
-    ):
-        self.span_id = span_id
-        self.operation_name = operation_name
-        self.master_duration = master_duration
-        self.branch_duration = branch_duration
-        self.improvement = f"{round((1 - branch_duration / master_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 = {
-            "span_id": self.span_id,
-            "operation_name": self.operation_name,
-            "master_duration": self.master_duration,
-            "branch_duration": self.branch_duration,
-            "improvement": self.improvement,
-        }
 
-        if include_children:
-            res["children"] = [child.to_dict() for child in self.children.values()]
+def match_children(
+    master_child: dict,
+    branch_child: dict,
+    master_children: list[dict],
+    branch_children: list[dict],
+):
+    # Attempt to find a matching child in Branch Children for the current Master Child
+    for i, branch_candidate in enumerate(branch_children):
+        if branch_candidate.get("operationName") == master_child.get("operationName"):
+            branch_children.pop(i)
+            return master_child, branch_candidate
 
-        return res
+    # Attempt to find a matching child in Master Children for the current Branch Child
+    for i, master_candidate in enumerate(master_children):
+        if master_candidate.get("operationName") == branch_child.get("operationName"):
+            master_children.pop(i)
+            return master_candidate, branch_child
 
+    return master_child, branch_child
 
-class DiffTree:
-    def __init__(self, master_trace: dict, branch_trace: dict):
-        self.root_span = self.__build_diff_tree(master_trace, branch_trace)
 
-    def __build_diff_tree(self, master_trace: dict, branch_trace: dict):
-        assert master_trace["operationName"] == branch_trace["operationName"]
+def compare_children(master_children: list[dict], branch_children: list[dict]):
+    result = []
+    master_children_sorted = sort_children_by_start_time(master_children)
+    branch_children_sorted = sort_children_by_start_time(branch_children)
 
-        diff_span = DiffSpan(
-            master_trace["spanID"],
-            master_trace["operationName"],
-            master_trace["duration"],
-            branch_trace["duration"],
-        )
+    while master_children_sorted or branch_children_sorted:
+        master_child = master_children_sorted.pop(0) if master_children_sorted else None
+        branch_child = branch_children_sorted.pop(0) if branch_children_sorted else None
 
-        # Recursively build the tree for child spans
-        for head_child_span, branch_child_span in zip(
-            master_trace["children"], branch_trace["children"]
+        if (
+            master_child
+            and branch_child
+            and master_child.get("operationName") != branch_child.get("operationName")
         ):
-            child_span = self.__build_diff_tree(head_child_span, branch_child_span)
-            diff_span.add_child(head_child_span["spanID"], child_span)
-
-        return diff_span
-
-    def to_dict(self):
-        return self.root_span.to_dict()
+            # Find the matching pair if they are out of order
+            master_child, branch_child = match_children(
+                master_child,
+                branch_child,
+                master_children_sorted,
+                branch_children_sorted,
+            )
+
+        result.append(compare_span(master_child, branch_child))
+
+    return result
+
+
+def compare_traces(master_trace: dict, branch_trace: dict):
+    assert master_trace.get("operationName") == branch_trace.get("operationName")
+    assert master_trace.get("tags") == branch_trace.get("tags")
+
+    result = {
+        "operationName": master_trace.get("operationName"),
+        "master_duration": master_trace["duration"],
+        "branch_duration": branch_trace["duration"],
+        "improvement": f"{round((1 - branch_trace['duration'] / master_trace['duration']) * 100)}%",
+        "tags": master_trace.get("tags"),
+        "children": compare_children(
+            master_trace["children"], branch_trace["children"]
+        ),
+    }
 
-    def to_json_file(self, operation: str):
-        current_dir = pathlib.Path(__file__).parent.absolute()
-        with open(f"{current_dir}/diff/{operation}.json", "w") as file:
-            file.write(json.dumps(self.to_dict(), indent=4))
+    return result
 
 
 def main():
     current_dir = pathlib.Path(__file__).parent.absolute()
 
+    root_span_diffs = {}
     for operation in ["pyg_to_arangodb", "arangodb_to_pyg"]:
         master_trace = json.load(open(f"{current_dir}/traces/master/{operation}.json"))
         branch_trace = json.load(open(f"{current_dir}/traces/branch/{operation}.json"))
+        diff_trace = compare_traces(master_trace, branch_trace)
+
+        with open(f"{current_dir}/diff/{operation}.json", "w") as file:
+            file.write(json.dumps(diff_trace, indent=4))
+
+        root_span_diffs[operation] = {
+            "master_duration": diff_trace["master_duration"],
+            "branch_duration": diff_trace["branch_duration"],
+            "improvement": diff_trace["improvement"],
+        }
 
-        diff_tree = DiffTree(master_trace, branch_trace)
-        diff_tree.to_json_file(operation)
+    print("-" * 50)
+    print(json.dumps(root_span_diffs, indent=4))
+    print("-" * 50)
 
-        print("-" * 50)
-        print(json.dumps(diff_tree.root_span.to_dict(include_children=False), indent=4))
-        print("-" * 50)
+    return root_span_diffs
 
 
 if __name__ == "__main__":