def disconnect_at_roots(current_dag: Dag, available: set, new_roots: list): previous_parents = set() create_op_lookup = dict() for root in new_roots: for parent in copy(root.parents): if parent in available: create_op = None if parent not in previous_parents: create_op = Create(deepcopy(parent.out_rel)) # create op is in same mode as root create_op.is_mpc = root.is_mpc previous_parents.add(parent) create_op_lookup[parent.out_rel.name] = create_op else: create_op = create_op_lookup[parent.out_rel.name] # unlink root from parent parent.children.remove(root) # insert create op between parent and root root.replace_parent(parent, create_op) # connect create op with root create_op.children.add(root) # keep track of parents we have already visited previous_parents.add(parent) create_op_lookup[create_op.out_rel.name] = create_op if root in current_dag.roots: current_dag.roots.remove(root) parent_roots = set().union(*[root.parents for root in new_roots]) for root in new_roots: if isinstance(root, Create): parent_roots.add(root) return OpDag(set(parent_roots)), available
def rewrite(self, dag: saldag.OpDag): """ Traverse topologically sorted DAG, inspect each node. """ ordered = dag.top_sort() if self.reverse: ordered = ordered[::-1] for node in ordered: print(type(self).__name__, "rewriting", node.out_rel.name) if isinstance(node, saldag.Aggregate): self._rewrite_aggregate(node) elif isinstance(node, saldag.Divide): self._rewrite_divide(node) elif isinstance(node, saldag.Project): self._rewrite_project(node) elif isinstance(node, saldag.Filter): self._rewrite_filter(node) elif isinstance(node, saldag.Multiply): self._rewrite_multiply(node) elif isinstance(node, saldag.RevealJoin): self._rewrite_reveal_join(node) elif isinstance(node, saldag.HybridJoin): self._rewrite_hybrid_join(node) elif isinstance(node, saldag.Join): self._rewrite_join(node) elif isinstance(node, saldag.Concat): self._rewrite_concat(node) elif isinstance(node, saldag.Close): self._rewrite_close(node) elif isinstance(node, saldag.Open): self._rewrite_open(node) elif isinstance(node, saldag.Create): self._rewrite_create(node) elif isinstance(node, saldag.Distinct): self._rewrite_distinct(node) else: msg = "Unknown class " + type(node).__name__ raise Exception(msg)
def _merge_dags(left_dag, right_dag): # TODO: should go inside dagutils, once dagutils exists # to merge, we only need to combine roots roots = left_dag.roots.union(right_dag.roots) return OpDag(roots)