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,21 @@ def update_chunk(cg: ChunkedGraph, chunk_coords: list[int]):
     nodes = []
     nodes_ts = []
     earliest_ts = cg.get_earliest_timestamp()
+    corrupt_nodes = []
+    corrupt_nodes_ts = []
     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)
+                corrupt_nodes_ts.append(ts)
 
     if len(nodes) > 0:
         logging.info(f"processing {len(nodes)} nodes.")
@@ -156,3 +161,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, corrupt_nodes_ts, 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,30 @@ def _update_cross_edges_helper(args):
     parents = cg.get_parents(nodes, fail_to_zero=True)
 
     tasks = []
+    corrupt_nodes = []
+    corrupt_nodes_ts = []
+    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)
+                corrupt_nodes_ts.append(node_ts)
+        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, corrupt_nodes_ts, CHILDREN)
 
 
 def update_chunk(
@@ -164,7 +177,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,38 @@ def get_parent_timestamps(
             ts = cell.timestamp
             result[k].add(earliest_ts if ts < earliest_ts else ts)
     return result
+
+
+def _fix_corrupt_node(cg: ChunkedGraph, node: int, children: np.ndarray):
+    """
+    Removes this 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)
+    children_d = cg.client.read_nodes(node_ids=children, properties=Hierarchy.Parent)
+    for child, parent_cells in children_d.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()
+
+
+def fix_corrupt_nodes(cg: ChunkedGraph, nodes: list, nodes_ts: list, children_d: dict):
+    _children_d = {node: children_d[node] for node in nodes}
+    end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, _children_d, layer=2)
+    for node, end_ts in zip(nodes, end_timestamps):
+        assert end_ts is None, f"{node}: {end_ts}"
+        _fix_corrupt_node(cg, node, _children_d[node])