feat: add set_max_node_id, tests

Author: akhileshh

kvdbclient/attributes.py

@@ -132,6 +132,13 @@ class Connectivity:
         serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID, shape=(-1, 2)),
     )
 
+    # new edges as a result of supervoxel split
+    SplitEdges = _Attribute(
+        key=b"split_edges",
+        family_id="4",
+        serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID, shape=(-1, 2)),
+    )
+
 
 class Hierarchy:
     Child = _Attribute(
@@ -166,6 +173,18 @@ class Hierarchy:
         key=b"stale_ts", family_id="3", serializer=serializers.Pickle()
     )
 
+    FormerIdentity = _Attribute(
+        key=b"former_ids",
+        family_id="0",
+        serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID),
+    )
+
+    NewIdentity = _Attribute(
+        key=b"new_ids",
+        family_id="0",
+        serializer=serializers.NumPyArray(dtype=basetypes.NODE_ID),
+    )
+
 
 class TableMeta:
     key = b"meta"

kvdbclient/base.py

@@ -541,6 +541,13 @@ def create_node_id(
     ) -> basetypes.NODE_ID:
         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
+    ) -> None:
+        """Set max segment ID for a given chunk."""
+        size = int(np.uint64(chunk_id) ^ np.uint64(node_id))
+        self._get_ids_range(serialize_uint64(chunk_id, counter=True), size)
+
     def get_max_node_id(
         self, chunk_id: basetypes.CHUNK_ID, root_chunk=False
     ) -> basetypes.NODE_ID:

tests/test_bigtable_client.py

@@ -400,6 +400,42 @@ def test_get_max_node_id(self, bt_client):
         max_id = bt_client.get_max_node_id(chunk_id)
         assert max_id > chunk_id
 
+    def test_set_max_node_id_fresh(self, bt_client):
+        """On a fresh counter, set_max_node_id sets the counter to the segment ID."""
+        chunk_id = np.uint64(1 << 32)
+        node_id = chunk_id | np.uint64(10)
+        bt_client.set_max_node_id(chunk_id, node_id)
+        max_id = bt_client.get_max_node_id(chunk_id)
+        assert max_id == node_id
+
+    def test_set_max_node_id_then_create_no_collision(self, bt_client):
+        """After set_max_node_id, create_node_id should return IDs above the set max."""
+        chunk_id = np.uint64(1 << 32)
+        node_id = chunk_id | np.uint64(10)
+        bt_client.set_max_node_id(chunk_id, node_id)
+        new_id = bt_client.create_node_id(chunk_id)
+        segment_id = int(np.uint64(new_id) ^ np.uint64(chunk_id))
+        assert segment_id > 10
+
+    def test_set_max_node_id_is_additive(self, bt_client):
+        """Calling set_max_node_id twice increments cumulatively (not idempotent)."""
+        chunk_id = np.uint64(1 << 32)
+        node_id = chunk_id | np.uint64(5)
+        bt_client.set_max_node_id(chunk_id, node_id)
+        bt_client.set_max_node_id(chunk_id, node_id)
+        max_id = bt_client.get_max_node_id(chunk_id)
+        # Counter was incremented by 5 twice -> segment_id is 10
+        assert int(np.uint64(max_id) ^ np.uint64(chunk_id)) == 10
+
+    def test_set_max_node_id_after_create(self, bt_client):
+        """set_max_node_id after create_node_ids advances counter further."""
+        chunk_id = np.uint64(1 << 32)
+        bt_client.create_node_ids(chunk_id, 3)  # counter at 3
+        node_id = chunk_id | np.uint64(7)
+        bt_client.set_max_node_id(chunk_id, node_id)  # increments by 7 -> counter at 10
+        max_id = bt_client.get_max_node_id(chunk_id)
+        assert int(np.uint64(max_id) ^ np.uint64(chunk_id)) == 10
+
 
 class TestOperationIds:
     def test_create_unique(self, bt_client):

tests/test_hbase_client.py

@@ -405,6 +405,42 @@ def test_create_batch(self, hbase_client):
         assert len(ids) == 5
         assert len(np.unique(ids)) == 5
 
+    def test_set_max_node_id_fresh(self, hbase_client):
+        """On a fresh counter, set_max_node_id sets the counter to the segment ID."""
+        chunk_id = np.uint64(1 << 32)
+        node_id = chunk_id | np.uint64(10)
+        hbase_client.set_max_node_id(chunk_id, node_id)
+        max_id = hbase_client.get_max_node_id(chunk_id)
+        assert max_id == node_id
+
+    def test_set_max_node_id_then_create_no_collision(self, hbase_client):
+        """After set_max_node_id, create_node_id should return IDs above the set max."""
+        chunk_id = np.uint64(1 << 32)
+        node_id = chunk_id | np.uint64(10)
+        hbase_client.set_max_node_id(chunk_id, node_id)
+        new_id = hbase_client.create_node_id(chunk_id)
+        segment_id = int(np.uint64(new_id) ^ np.uint64(chunk_id))
+        assert segment_id > 10
+
+    def test_set_max_node_id_is_additive(self, hbase_client):
+        """Calling set_max_node_id twice increments cumulatively (not idempotent)."""
+        chunk_id = np.uint64(1 << 32)
+        node_id = chunk_id | np.uint64(5)
+        hbase_client.set_max_node_id(chunk_id, node_id)
+        hbase_client.set_max_node_id(chunk_id, node_id)
+        max_id = hbase_client.get_max_node_id(chunk_id)
+        # Counter was incremented by 5 twice -> segment_id is 10
+        assert int(np.uint64(max_id) ^ np.uint64(chunk_id)) == 10
+
+    def test_set_max_node_id_after_create(self, hbase_client):
+        """set_max_node_id after create_node_ids advances counter further."""
+        chunk_id = np.uint64(1 << 32)
+        hbase_client.create_node_ids(chunk_id, 3)  # counter at 3
+        node_id = chunk_id | np.uint64(7)
+        hbase_client.set_max_node_id(chunk_id, node_id)  # increments by 7 -> counter at 10
+        max_id = hbase_client.get_max_node_id(chunk_id)
+        assert int(np.uint64(max_id) ^ np.uint64(chunk_id)) == 10
+
 
 class TestHBaseOperationIds:
     def test_create_unique(self, hbase_client):