fix(migration): erase corrupt ids from failed edits

Author: akhileshh

pychunkedgraph/ingest/upgrade/atomic_layer.py

@@ -11,7 +11,7 @@
 from pychunkedgraph.graph.attributes import Connectivity, Hierarchy
 from pychunkedgraph.graph.utils import serializers
 
-from .utils import get_end_timestamps, get_parent_timestamps
+from .utils import fix_corrupt_nodes, get_end_timestamps, get_parent_timestamps
 
 CHILDREN = {}
 
@@ -136,16 +136,19 @@ def update_chunk(cg: ChunkedGraph, chunk_coords: list[int]):
     nodes = []
     nodes_ts = []
     earliest_ts = cg.get_earliest_timestamp()
+    corrupt_nodes = []
     for k, v in rr.items():
         try:
-            _ = v[Hierarchy.Parent]
-            nodes.append(k)
             CHILDREN[k] = v[Hierarchy.Child][0].value
             ts = v[Hierarchy.Child][0].timestamp
+            _ = v[Hierarchy.Parent]
+            nodes.append(k)
             nodes_ts.append(earliest_ts if ts < earliest_ts else ts)
         except KeyError:
-            # invalid nodes from failed tasks w/o parent column entry
-            continue
+            # ignore invalid nodes from failed ingest tasks, w/o parent column entry
+            # retain invalid nodes from edits to fix the hierarchy
+            if ts > earliest_ts:
+                corrupt_nodes.append(k)
 
     if len(nodes) > 0:
         logging.info(f"processing {len(nodes)} nodes.")
@@ -156,3 +159,7 @@ def update_chunk(cg: ChunkedGraph, chunk_coords: list[int]):
     rows = update_nodes(cg, nodes, nodes_ts)
     cg.client.write(rows)
     logging.info(f"mutations: {len(rows)}, time: {time.time() - start}")
+
+    if len(corrupt_nodes) > 0:
+        logging.info(f"found {len(corrupt_nodes)} corrupt nodes {corrupt_nodes[:3]}...")
+        fix_corrupt_nodes(cg, corrupt_nodes, CHILDREN)

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -17,7 +17,7 @@
 from pychunkedgraph.graph.types import empty_2d
 from pychunkedgraph.utils.general import chunked
 
-from .utils import get_end_timestamps, get_parent_timestamps
+from .utils import fix_corrupt_nodes, get_end_timestamps, get_parent_timestamps
 
 
 CHILDREN = {}
@@ -145,17 +145,28 @@ def _update_cross_edges_helper(args):
     parents = cg.get_parents(nodes, fail_to_zero=True)
 
     tasks = []
+    corrupt_nodes = []
+    earliest_ts = cg.get_earliest_timestamp()
     for node, parent, node_ts in zip(nodes, parents, nodes_ts):
         if parent == 0:
-            # invalid id caused by failed ingest task / edits
-            continue
-        tasks.append((cg, layer, node, node_ts))
+            # ignore invalid nodes from failed ingest tasks, w/o parent column entry
+            # retain invalid nodes from edits to fix the hierarchy
+            if node_ts > earliest_ts:
+                corrupt_nodes.append(node)
+        else:
+            tasks.append((cg, layer, node, node_ts))
 
     with ThreadPoolExecutor(max_workers=4) as executor:
-        futures = [executor.submit(_update_cross_edges_helper_thread, task) for task in tasks]
+        futures = [
+            executor.submit(_update_cross_edges_helper_thread, task) for task in tasks
+        ]
         for future in tqdm(as_completed(futures), total=len(futures)):
             rows.extend(future.result())
+
     cg.client.write(rows)
+    if len(corrupt_nodes) > 0:
+        logging.info(f"found {len(corrupt_nodes)} corrupt nodes {corrupt_nodes[:3]}...")
+        fix_corrupt_nodes(cg, corrupt_nodes, CHILDREN)
 
 
 def update_chunk(
@@ -164,7 +175,7 @@ def update_chunk(
     """
     Iterate over all layer IDs in a chunk and update their cross chunk edges.
     """
-    debug =  nodes is not None
+    debug = nodes is not None
     start = time.time()
     x, y, z = chunk_coords
     chunk_id = cg.get_chunk_id(layer=layer, x=x, y=y, z=z)

pychunkedgraph/ingest/upgrade/utils.py

@@ -1,11 +1,13 @@
 # pylint: disable=invalid-name, missing-docstring
 
 from collections import defaultdict
-from datetime import datetime, timezone
+from datetime import datetime, timedelta
 
 import numpy as np
 from pychunkedgraph.graph import ChunkedGraph
 from pychunkedgraph.graph.attributes import Hierarchy
+from pychunkedgraph.graph.utils import serializers
+from google.cloud.bigtable.row_filters import TimestampRange
 
 
 def exists_as_parent(cg: ChunkedGraph, parent, nodes) -> bool:
@@ -102,3 +104,34 @@ def get_parent_timestamps(
             ts = cell.timestamp
             result[k].add(earliest_ts if ts < earliest_ts else ts)
     return result
+
+
+def fix_corrupt_nodes(cg: ChunkedGraph, nodes: list, children_d: dict):
+    """
+    Iteratively removes a node from parent column of its children.
+    Then removes the node iteself, effectively erasing it.
+    """
+    table = cg.client._table
+    batcher = table.mutations_batcher(flush_count=500)
+    for node in nodes:
+        children = children_d[node]
+        _map = cg.client.read_nodes(node_ids=children, properties=Hierarchy.Parent)
+
+        for child, parent_cells in _map.items():
+            row = table.direct_row(serializers.serialize_uint64(child))
+            for cell in parent_cells:
+                if cell.value == node:
+                    start = cell.timestamp
+                    end = start + timedelta(microseconds=1)
+                    row.delete_cell(
+                        column_family_id=Hierarchy.Parent.family_id,
+                        column=Hierarchy.Parent.key,
+                        time_range=TimestampRange(start=start, end=end),
+                    )
+                    batcher.mutate(row)
+
+        row = table.direct_row(serializers.serialize_uint64(node))
+        row.delete()
+        batcher.mutate(row)
+
+    batcher.flush()