hatchet/graphframe.py
def squash(self):
        """Rewrite the Graph to include only nodes present in the DataFrame's rows.

        This can be used to simplify the Graph, or to normalize Graph
        indexes between two GraphFrames.
        """
        index_names = self.dataframe.index.names
        self.dataframe.reset_index(inplace=True)

        # create new nodes for each unique node in the old dataframe
        old_to_new = {n: n.copy() for n in set(self.dataframe["node"])}
        for i in old_to_new:
            old_to_new[i]._hatchet_nid = i._hatchet_nid

        # Maintain sets of connections to make for each old node.
        # Start with old -> new mapping and update as we traverse subgraphs.
        connections = defaultdict(lambda: set())
        connections.update({k: {v} for k, v in old_to_new.items()})

        new_roots = []  # list of new roots

        # connect new nodes to children according to transitive
        # relationships in the old graph.
        def rewire(node, new_parent, visited):
            # make all transitive connections for the node we're visiting
            for n in connections[node]:
                if new_parent:
                    # there is a parent in the new graph; connect it
                    if n not in new_parent.children:
                        new_parent.add_child(n)
                        n.add_parent(new_parent)

                elif n not in new_roots:
                    # this is a new root
                    new_roots.append(n)

            new_node = old_to_new.get(node)
            transitive = set()
            if node not in visited:
                visited.add(node)
                for child in node.children:
                    transitive |= rewire(child, new_node or new_parent, visited)

            if new_node:
                # since new_node exists in the squashed graph, we only
                # need to connect new_node
                return {new_node}
            else:
                # connect parents to the first transitively reachable
                # new_nodes of nodes we're removing with this squash
                connections[node] |= transitive
                return connections[node]

        # run rewire for each root and make a new graph
        visited = set()
        for root in self.graph.roots:
            rewire(root, None, visited)
        graph = Graph(new_roots)
        graph.enumerate_traverse()

        # reindex new dataframe with new nodes
        df = self.dataframe.copy()
        df["node"] = df["node"].apply(lambda x: old_to_new[x])

        # at this point, the graph is potentially invalid, as some nodes
        # may have children with identical frames.
        merges = graph.normalize()
        df["node"] = df["node"].apply(lambda n: merges.get(n, n))

        self.dataframe.set_index(index_names, inplace=True)
        df.set_index(index_names, inplace=True)
        # create dict that stores aggregation function for each column
        agg_dict = {}
        for col in df.columns.tolist():
            if col in self.exc_metrics + self.inc_metrics:
                agg_dict[col] = np.sum
            else:
                agg_dict[col] = lambda x: x.iloc[0]

        # perform a groupby to merge nodes with the same callpath
        agg_df = df.groupby(index_names).agg(agg_dict)
        agg_df.sort_index(inplace=True)

        # put it all together
        new_gf = GraphFrame(graph, agg_df, self.exc_metrics, self.inc_metrics)
        new_gf.update_inclusive_columns()
        return new_gf
Rewrite the Graph to include only nodes present in the DataFrame's rows.

        This can be used to simplify the Graph, or to normalize Graph
        indexes between two GraphFrames.
        
Rewrite the Graph to include only nodes present in the DataFrame's rows.
This can be used to simplify the Graph, or to normalize Graph
indexes between two GraphFrames.
Rewrite the Graph to include only nodes present in the DataFrame ' s rows . This can be used to simplify the Graph or to normalize Graph indexes between two GraphFrames .
def squash(self):
        index_names = self.dataframe.index.names
        self.dataframe.reset_index(inplace=True)
        old_to_new = {n: n.copy() for n in set(self.dataframe["node"])}
        for i in old_to_new:
            old_to_new[i]._hatchet_nid = i._hatchet_nid
        connections = defaultdict(lambda: set())
        connections.update({k: {v} for k, v in old_to_new.items()})
        new_roots = []  
        def rewire(node, new_parent, visited):
            for n in connections[node]:
                if new_parent:
                    if n not in new_parent.children:
                        new_parent.add_child(n)
                        n.add_parent(new_parent)
                elif n not in new_roots:
                    new_roots.append(n)
            new_node = old_to_new.get(node)
            transitive = set()
            if node not in visited:
                visited.add(node)
                for child in node.children:
                    transitive |= rewire(child, new_node or new_parent, visited)
            if new_node:
                return {new_node}
            else:
                connections[node] |= transitive
                return connections[node]
        visited = set()
        for root in self.graph.roots:
            rewire(root, None, visited)
        graph = Graph(new_roots)
        graph.enumerate_traverse()
        df = self.dataframe.copy()
        df["node"] = df["node"].apply(lambda x: old_to_new[x])
        merges = graph.normalize()
        df["node"] = df["node"].apply(lambda n: merges.get(n, n))
        self.dataframe.set_index(index_names, inplace=True)
        df.set_index(index_names, inplace=True)
        agg_dict = {}
        for col in df.columns.tolist():
            if col in self.exc_metrics + self.inc_metrics:
                agg_dict[col] = np.sum
            else:
                agg_dict[col] = lambda x: x.iloc[0]
        agg_df = df.groupby(index_names).agg(agg_dict)
        agg_df.sort_index(inplace=True)
        new_gf = GraphFrame(graph, agg_df, self.exc_metrics, self.inc_metrics)
        new_gf.update_inclusive_columns()
        return new_gf
def squash ( self ) : index_names = self . dataframe . index . names self . dataframe . reset_index ( inplace = True ) old_to_new = { n : n . copy ( ) for n in set ( self . dataframe [ "node" ] ) } for i in old_to_new : old_to_new [ i ] . _hatchet_nid = i . _hatchet_nid connections = defaultdict ( lambda : set ( ) ) connections . update ( { k : { v } for k , v in old_to_new . items ( ) } ) new_roots = [ ] def rewire ( node , new_parent , visited ) : for n in connections [ node ] : if new_parent : if n not in new_parent . children : new_parent . add_child ( n ) n . add_parent ( new_parent ) elif n not in new_roots : new_roots . append ( n ) new_node = old_to_new . get ( node ) transitive = set ( ) if node not in visited : visited . add ( node ) for child in node . children : transitive |= rewire ( child , new_node or new_parent , visited ) if new_node : return { new_node } else : connections [ node ] |= transitive return connections [ node ] visited = set ( ) for root in self . graph . roots : rewire ( root , None , visited ) graph = Graph ( new_roots ) graph . enumerate_traverse ( ) df = self . dataframe . copy ( ) df [ "node" ] = df [ "node" ] . apply ( lambda x : old_to_new [ x ] ) merges = graph . normalize ( ) df [ "node" ] = df [ "node" ] . apply ( lambda n : merges . get ( n , n ) ) self . dataframe . set_index ( index_names , inplace = True ) df . set_index ( index_names , inplace = True ) agg_dict = { } for col in df . columns . tolist ( ) : if col in self . exc_metrics + self . inc_metrics : agg_dict [ col ] = np . sum else : agg_dict [ col ] = lambda x : x . iloc [ 0 ] agg_df = df . groupby ( index_names ) . agg ( agg_dict ) agg_df . sort_index ( inplace = True ) new_gf = GraphFrame ( graph , agg_df , self . exc_metrics , self . inc_metrics ) new_gf . update_inclusive_columns ( ) return new_gf
Rewrite the Graph to include only nodes present in the DataFrame's rows.
Rewrite the Graph to include only nodes present in the DataFrame ' s rows .
"""Rewrite the Graph to include only nodes present in the DataFrame's rows.

        This can be used to simplify the Graph, or to normalize Graph
        indexes between two GraphFrames.
        """ # create new nodes for each unique node in the old dataframe # Maintain sets of connections to make for each old node. # Start with old -> new mapping and update as we traverse subgraphs. # list of new roots # connect new nodes to children according to transitive # relationships in the old graph. # make all transitive connections for the node we're visiting # there is a parent in the new graph; connect it # this is a new root # since new_node exists in the squashed graph, we only # need to connect new_node # connect parents to the first transitively reachable # new_nodes of nodes we're removing with this squash # run rewire for each root and make a new graph # reindex new dataframe with new nodes # at this point, the graph is potentially invalid, as some nodes # may have children with identical frames. # create dict that stores aggregation function for each column # perform a groupby to merge nodes with the same callpath # put it all together
self