@@ -494,7 +494,7 @@ def __init__(
494494 self ._old_hierarchy_d = old_hierarchy_d
495495 self ._new_old_id_d = new_old_id_d
496496 self ._old_new_id_d = old_new_id_d
497- self ._new_ids_d = defaultdict (list ) # new IDs in each layer
497+ self ._new_ids_d = defaultdict (list )
498498 self ._operation_id = operation_id
499499 self ._time_stamp = time_stamp
500500 self ._last_successful_ts = parent_ts
@@ -572,7 +572,6 @@ def _update_cross_edge_cache(self, parent, children):
572572 time_stamp = self ._last_successful_ts ,
573573 )
574574 edge_parents_d = dict (zip (edge_nodes , edge_parents ))
575-
576575 new_cx_edges_d = {}
577576 for layer in range (parent_layer , self .cg .meta .layer_count ):
578577 edges = cx_edges_d .get (layer , types .empty_2d )
@@ -583,37 +582,44 @@ def _update_cross_edge_cache(self, parent, children):
583582 assert np .all (edges [:, 0 ] == parent )
584583 self .cg .cache .cross_chunk_edges_cache [parent ] = new_cx_edges_d
585584
586- def _update_neighbor_parents (self , neighbor , ceil_layer , updated ) -> list :
587- updated_parents = []
585+ def _update_neighbor_parents (self , neighbor , ceil_layer : int , updated : set ) -> list :
586+ """helper for `_update_skipped_neighbors`"""
587+ parents = []
588588 while True :
589589 parent = self .cg .get_parent (neighbor , time_stamp = self ._last_successful_ts )
590590 parent_layer = self .cg .get_chunk_layer (parent )
591591 if parent_layer >= ceil_layer or parent in updated :
592592 break
593593 children = self .cg .get_children (parent )
594594 self ._update_cross_edge_cache (parent , children )
595- updated_parents .append (parent )
595+ parents .append (parent )
596596 neighbor = parent
597- return updated_parents
597+ return parents
598598
599599 def _update_skipped_neighbors (self , node , layer , parent_layer ):
600+ """
601+ Updates cross edges of neighbors of a skip connection node.
602+ Neighbors of such nodes can have parents at contiguous layers.
603+
604+ This method updates cross edges of all such parents
605+ from `layer` through `parent_layer`.
606+ """
600607 updated_parents = set ()
601608 cx_edges_d = self .cg .cache .cross_chunk_edges_cache [node ]
602- for l in range (layer , parent_layer + 1 ):
603- layer_edges = cx_edges_d .get (l , types .empty_2d )
609+ for _layer in range (layer , parent_layer + 1 ):
610+ layer_edges = cx_edges_d .get (_layer , types .empty_2d )
604611 neighbors = layer_edges [:, 1 ]
605612 for n in neighbors :
606613 if n in self ._new_old_id_d :
607614 # ignore new ids
608615 continue
609616 res = self ._update_neighbor_parents (n , parent_layer , updated_parents )
610617 updated_parents .update (res )
611-
612618 updated_entries = []
613619 for parent in updated_parents :
614620 val_dict = {}
615- for layer , edges in self .cg .cache .cross_chunk_edges_cache [parent ].items ():
616- val_dict [attributes .Connectivity .CrossChunkEdge [layer ]] = edges
621+ for _layer , edges in self .cg .cache .cross_chunk_edges_cache [parent ].items ():
622+ val_dict [attributes .Connectivity .CrossChunkEdge [_layer ]] = edges
617623 rkey = serialize_uint64 (parent )
618624 row = self .cg .client .mutate_row (rkey , val_dict , time_stamp = self ._time_stamp )
619625 updated_entries .append (row )
0 commit comments