| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211 |
- import re
- class ProfileStats:
- """
- ProfileStats, runtime execution statistics of operation.
- """
- def __init__(self, records_produced, execution_time):
- self.records_produced = records_produced
- self.execution_time = execution_time
- class Operation:
- """
- Operation, single operation within execution plan.
- """
- def __init__(self, name, args=None, profile_stats=None):
- """
- Create a new operation.
- Args:
- name: string that represents the name of the operation
- args: operation arguments
- profile_stats: profile statistics
- """
- self.name = name
- self.args = args
- self.profile_stats = profile_stats
- self.children = []
- def append_child(self, child):
- if not isinstance(child, Operation) or self is child:
- raise Exception("child must be Operation")
- self.children.append(child)
- return self
- def child_count(self):
- return len(self.children)
- def __eq__(self, o: object) -> bool:
- if not isinstance(o, Operation):
- return False
- return self.name == o.name and self.args == o.args
- def __str__(self) -> str:
- args_str = "" if self.args is None else " | " + self.args
- return f"{self.name}{args_str}"
- class ExecutionPlan:
- """
- ExecutionPlan, collection of operations.
- """
- def __init__(self, plan):
- """
- Create a new execution plan.
- Args:
- plan: array of strings that represents the collection operations
- the output from GRAPH.EXPLAIN
- """
- if not isinstance(plan, list):
- raise Exception("plan must be an array")
- if isinstance(plan[0], bytes):
- plan = [b.decode() for b in plan]
- self.plan = plan
- self.structured_plan = self._operation_tree()
- def _compare_operations(self, root_a, root_b):
- """
- Compare execution plan operation tree
- Return: True if operation trees are equal, False otherwise
- """
- # compare current root
- if root_a != root_b:
- return False
- # make sure root have the same number of children
- if root_a.child_count() != root_b.child_count():
- return False
- # recursively compare children
- for i in range(root_a.child_count()):
- if not self._compare_operations(root_a.children[i], root_b.children[i]):
- return False
- return True
- def __str__(self) -> str:
- def aggraget_str(str_children):
- return "\n".join(
- [
- " " + line
- for str_child in str_children
- for line in str_child.splitlines()
- ]
- )
- def combine_str(x, y):
- return f"{x}\n{y}"
- return self._operation_traverse(
- self.structured_plan, str, aggraget_str, combine_str
- )
- def __eq__(self, o: object) -> bool:
- """Compares two execution plans
- Return: True if the two plans are equal False otherwise
- """
- # make sure 'o' is an execution-plan
- if not isinstance(o, ExecutionPlan):
- return False
- # get root for both plans
- root_a = self.structured_plan
- root_b = o.structured_plan
- # compare execution trees
- return self._compare_operations(root_a, root_b)
- def _operation_traverse(self, op, op_f, aggregate_f, combine_f):
- """
- Traverse operation tree recursively applying functions
- Args:
- op: operation to traverse
- op_f: function applied for each operation
- aggregate_f: aggregation function applied for all children of a single operation
- combine_f: combine function applied for the operation result and the children result
- """ # noqa
- # apply op_f for each operation
- op_res = op_f(op)
- if len(op.children) == 0:
- return op_res # no children return
- else:
- # apply _operation_traverse recursively
- children = [
- self._operation_traverse(child, op_f, aggregate_f, combine_f)
- for child in op.children
- ]
- # combine the operation result with the children aggregated result
- return combine_f(op_res, aggregate_f(children))
- def _operation_tree(self):
- """Build the operation tree from the string representation"""
- # initial state
- i = 0
- level = 0
- stack = []
- current = None
- def _create_operation(args):
- profile_stats = None
- name = args[0].strip()
- args.pop(0)
- if len(args) > 0 and "Records produced" in args[-1]:
- records_produced = int(
- re.search("Records produced: (\\d+)", args[-1]).group(1)
- )
- execution_time = float(
- re.search("Execution time: (\\d+.\\d+) ms", args[-1]).group(1)
- )
- profile_stats = ProfileStats(records_produced, execution_time)
- args.pop(-1)
- return Operation(
- name, None if len(args) == 0 else args[0].strip(), profile_stats
- )
- # iterate plan operations
- while i < len(self.plan):
- current_op = self.plan[i]
- op_level = current_op.count(" ")
- if op_level == level:
- # if the operation level equal to the current level
- # set the current operation and move next
- child = _create_operation(current_op.split("|"))
- if current:
- current = stack.pop()
- current.append_child(child)
- current = child
- i += 1
- elif op_level == level + 1:
- # if the operation is child of the current operation
- # add it as child and set as current operation
- child = _create_operation(current_op.split("|"))
- current.append_child(child)
- stack.append(current)
- current = child
- level += 1
- i += 1
- elif op_level < level:
- # if the operation is not child of current operation
- # go back to it's parent operation
- levels_back = level - op_level + 1
- for _ in range(levels_back):
- current = stack.pop()
- level -= levels_back
- else:
- raise Exception("corrupted plan")
- return stack[0]
|