feat: set max node id func, to update L1 chunk max seg ID

Author: akhileshh

pychunkedgraph/graph/client/base.py

@@ -122,6 +122,10 @@ def create_node_ids(self, chunk_id):
     def create_node_id(self, chunk_id):
         """Generate a unique ID in the chunk."""
 
+    @abstractmethod
+    def set_max_node_id(self, chunk_id, node_id):
+        """Gets the current maximum node ID in the chunk."""
+
     @abstractmethod
     def get_max_node_id(self, chunk_id):
         """Gets the current maximum node ID in the chunk."""

pychunkedgraph/graph/client/bigtable/client.py

@@ -588,6 +588,19 @@ def create_node_id(
         """Generate a unique node ID in the chunk."""
         return self.create_node_ids(chunk_id, 1, root_chunk=root_chunk)[0]
 
+
+    def set_max_node_id(
+        self, chunk_id: np.uint64, node_id: np.uint64
+    ) -> basetypes.NODE_ID:
+        """Set max segment ID for a given chunk."""
+        size = int(np.uint64(chunk_id) ^ np.uint64(node_id))
+        key = serialize_uint64(chunk_id, counter=True)
+        column = attributes.Concurrency.Counter
+        row = self._table.append_row(key)
+        row.increment_cell_value(column.family_id, column.key, size)
+        row = row.commit()
+
+
     def get_max_node_id(
         self, chunk_id: basetypes.CHUNK_ID, root_chunk=False
     ) -> basetypes.NODE_ID:

pychunkedgraph/ingest/create/atomic_layer.py

@@ -35,7 +35,10 @@ def add_atomic_edges(
         return
 
     chunk_ids = cg.get_chunk_ids_from_node_ids(chunk_node_ids)
-    assert len(np.unique(chunk_ids)) == 1
+    assert len(np.unique(chunk_ids)) == 1, np.unique(chunk_ids)
+
+    max_node_id = np.max(chunk_node_ids)
+    cg.id_client.set_max_node_id(chunk_ids[0], max_node_id)
 
     graph, _, _, unique_ids = build_gt_graph(chunk_edge_ids, make_directed=True)
     ccs = connected_components(graph)

pychunkedgraph/repair/fake_edges.py

@@ -9,9 +9,9 @@
 from os import environ
 from typing import Optional
 
-environ["BIGTABLE_PROJECT"] = "<>"
-environ["BIGTABLE_INSTANCE"] = "<>"
-environ["GOOGLE_APPLICATION_CREDENTIALS"] = "<path>"
+# environ["BIGTABLE_PROJECT"] = "<>"
+# environ["BIGTABLE_INSTANCE"] = "<>"
+# environ["GOOGLE_APPLICATION_CREDENTIALS"] = "<path>"
 
 from pychunkedgraph.graph import edits
 from pychunkedgraph.graph import ChunkedGraph

pychunkedgraph/tests/test_uncategorized.py

@@ -139,7 +139,7 @@ def test_build_single_node(self, gen_graph):
         assert len(children) == 1 and children[0] == to_label(cg, 1, 0, 0, 0, 0)
         # Make sure there are not any more entries in the table
         # include counters, meta and version rows
-        assert len(res.rows) == 1 + 1 + 1 + 1 + 1
+        assert len(res.rows) == 1 + 1 + 1 + 1 + 1 + 1
 
     @pytest.mark.timeout(30)
     def test_build_single_edge(self, gen_graph):
@@ -192,7 +192,7 @@ def test_build_single_edge(self, gen_graph):
 
         # Make sure there are not any more entries in the table
         # include counters, meta and version rows
-        assert len(res.rows) == 2 + 1 + 1 + 1 + 1
+        assert len(res.rows) == 2 + 1 + 1 + 1 + 1 + 1
 
     @pytest.mark.timeout(30)
     def test_build_single_across_edge(self, gen_graph):
@@ -294,7 +294,7 @@ def test_build_single_across_edge(self, gen_graph):
 
         # Make sure there are not any more entries in the table
         # include counters, meta and version rows
-        assert len(res.rows) == 2 + 2 + 1 + 3 + 1 + 1
+        assert len(res.rows) == 2 + 2 + 1 + 3 + 1 + 1 + 2
 
     @pytest.mark.timeout(30)
     def test_build_single_edge_and_single_across_edge(self, gen_graph):
@@ -402,7 +402,7 @@ def test_build_single_edge_and_single_across_edge(self, gen_graph):
 
         # Make sure there are not any more entries in the table
         # include counters, meta and version rows
-        assert len(res.rows) == 3 + 2 + 1 + 3 + 1 + 1
+        assert len(res.rows) == 3 + 2 + 1 + 3 + 1 + 1 + 2
 
     @pytest.mark.timeout(120)
     def test_build_big_graph(self, gen_graph):