perf(merge): ⚡ segment into the label group and relabel it in place

Takes the pipeline from three full passes over the label volume to two. The
segment jobs write their tile-local ids straight into
image.zarr/labels/<name>/0, and the merge rewrites that array where it
stands instead of filling a scratch stage store and copying it across.

Relabelling in place is safe because the boundary scan completes before any
chunk is rewritten, so nothing still needs the original ids, and each worker
owns a disjoint, chunk-aligned region. zarr_native_merge detects that its
output is its input and skips creating the array -- creating it would delete
the very data it is about to read.

It only applies when the tile already fits the label chunk cap, because in
place the chunking cannot be changed and level 0 has to stay pageable for a
viewer. An oversized tile keeps the scratch store, and prepare records which
route it took in tiles.json so segment and merge agree.

The label group is not registered in labels/.zattrs until the merge finishes,
so a run that dies midway leaves an unregistered group that NGFF readers do
not list, and a rerun recreates it. That is the trade: the stage store used
to be a clean redo point.

Verified against the two-store path on a cross-boundary case (identical
labelling) and end to end, including a concurrent multi-config run.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: lguerard

src/patchworks/_merge.py

@@ -622,6 +622,21 @@ def zarr_native_merge(
             "zarr_native_merge: renumbered to 1..%d in the same LUT", n_objects
         )
 
+    # Relabelling straight back into the source array saves writing the whole
+    # volume a second time. It is safe because the boundary scan has already
+    # finished by this point, so nothing still needs the original ids, and
+    # each worker owns a disjoint, chunk-aligned region.
+    in_place = (staged_path, staged_component) == (out_path, out_component)
+    if in_place:
+        if output_chunks is not None and tuple(output_chunks) != tuple(
+            chunk_shape
+        ):
+            raise ValueError(
+                "output_chunks cannot differ from the source chunking when "
+                "merging in place (the array is rewritten, not recreated)"
+            )
+        logger.info("zarr_native_merge: relabeling in place, no second store")
+
     out_root = zarr.open_group(out_path, mode="a")
     out_chunks = tuple(chunk_shape)
     if output_chunks is not None:
@@ -636,11 +651,14 @@ def zarr_native_merge(
                 "output_chunks must divide the staged chunk shape so workers "
                 f"write whole chunks; axis/staged/output mismatches: {bad}"
             )
-    # Match the staged dtype: ids are already compact (dense by construction,
-    # and compacted again above when sequential), so nothing needs a wider one.
-    _create_zarr_label_array(
-        out_root, out_component, shape, out_chunks, dtype=arr.dtype
-    )
+    if not in_place:
+        # Match the staged dtype: ids are already compact (dense by
+        # construction, and compacted again above when sequential), so nothing
+        # needs a wider one. Creating this in place would delete the very
+        # array we are about to read.
+        _create_zarr_label_array(
+            out_root, out_component, shape, out_chunks, dtype=arr.dtype
+        )
 
     # Row-major, matching spatial_tiles' order -- so chunk i is tile i and the
     # offsets line up with the per-tile counts.
@@ -669,7 +687,7 @@ def zarr_native_merge(
     n_w = max(1, min(n_workers, max(1, len(tasks))))
     logger.info(
         "zarr_native_merge: relabeling %d chunks with %d worker(s)…",
-        n_chunks,
+        len(tasks),
         n_w,
     )
 

tests/test_distributed.py

@@ -192,6 +192,59 @@ def test_empty_chunks_are_skipped_not_just_short_circuited(tmp_path):
     assert written == ["0"], f"only the occupied chunk should exist: {written}"
 
 
+def test_in_place_merge_matches_the_two_store_merge(tmp_path):
+    """Relabelling back into the source must give the same labelling.
+
+    It is safe because the boundary scan finishes before any chunk is
+    rewritten, so nothing still needs the original ids -- and it saves writing
+    the whole volume a second time.
+    """
+    img = np.zeros((8, 48), "uint16")
+    img[1:4, 4:12] = 500  # inside tile 0
+    img[2:6, 14:26] = 500  # straddles the x=16 boundary
+    tile = (8, 16)
+
+    results = {}
+    for name in ("two_store", "in_place"):
+        stage = str(tmp_path / f"stage_{name}.zarr")
+        create_stage(stage, img.shape, tile)
+        counts = {
+            i: stage_tile(img, _fn, stage, i, tile_shape=tile, overlap=2)
+            for i in range(len(spatial_tiles(img.shape, tile)))
+        }
+        # in place = the merge's output *is* its input
+        out = (
+            stage if name == "in_place" else str(tmp_path / f"out_{name}.zarr")
+        )
+        component = "staged" if name == "in_place" else "labels"
+        results[name] = merge_tile_labels(
+            stage,
+            write_to=out,
+            input_component="staged",
+            output_component=component,
+            sequential_labels=True,
+            label_counts=counts,
+        ).compute()
+
+    assert np.array_equal(results["two_store"], results["in_place"]), (
+        "in-place relabelling changed the result"
+    )
+    assert set(np.unique(results["in_place"]).tolist()) == {0, 1, 2}
+
+    # And it must refuse a rechunk it cannot honour, rather than silently
+    # ignoring it: in place the array is rewritten, never recreated.
+    stage = str(tmp_path / "stage_guard.zarr")
+    create_stage(stage, img.shape, tile)
+    with pytest.raises(ValueError, match="cannot differ"):
+        merge_tile_labels(
+            stage,
+            write_to=stage,
+            input_component="staged",
+            output_component="staged",
+            output_chunks=(8, 8),
+        )
+
+
 def test_stage_tile_returns_dense_label_count(tmp_path):
     """stage_tile reports how many labels it wrote, and writes exactly 1..n."""
     img = np.zeros((8, 16), "uint16")

workflow/rules/merge.smk

@@ -2,7 +2,8 @@
 
 rule merge:
     input:
-        batch_done,
+        markers=batch_done,
+        tiles=TILES,
     output:
         touch(f"{RUN}/labels.done"),
     log:

workflow/scripts/merge.py

@@ -20,7 +20,7 @@
 from patchworks._chunks import _get_available_memory
 from patchworks.plugins.ome_zarr import register_labels
 
-from _pw import stage_path, start_log
+from _pw import load_tiles_json, stage_path, start_log
 
 start_log(snakemake.log[0])  # noqa: F821
 cfg = snakemake.config  # noqa: F821
@@ -29,7 +29,14 @@
 image_store = str(Path(work_dir) / "image.zarr")
 label_group = f"{image_store}/labels/{label_name}"
 
-staged = zarr.open_group(stage_path(work_dir, label_name), mode="r")["staged"]
+# prepare recorded where the segment jobs wrote: the label group's level 0
+# (merged in place) or a scratch stage store.
+manifest = load_tiles_json(snakemake.input.tiles)  # noqa: F821
+target_path = manifest["target_path"]
+target_component = manifest.get("target_component", "staged")
+in_place = bool(manifest.get("in_place", False))
+
+staged = zarr.open_group(target_path, mode="r")[target_component]
 
 # Size the relabel pool against what this job was actually granted, not the
 # node. Each worker holds roughly a few copies of one chunk, so the RAM budget
@@ -49,29 +56,30 @@
 # counts in lets the merge compute global id ranges by a cumulative sum,
 # replacing a full read+write of the store that existed only to renumber it.
 label_counts = {}
-for marker in snakemake.input:  # noqa: F821
+for marker in snakemake.input.markers:  # noqa: F821
     for index, n in json.loads(Path(marker).read_text())["counts"].items():
         label_counts[int(index)] = int(n)
 
-# Merge straight into the label group's level 0. Writing to a scratch
-# _merged.zarr and letting write_labels copy it across cost a full extra
-# read+write of the volume plus the scratch store's disk; register_labels
-# already expects level 0 to exist and only adds the pyramid and metadata.
-root = zarr.open_group(image_store, mode="a")
-parent = root.require_group("labels")
-if label_name in parent:
-    del parent[label_name]
-parent.require_group(label_name)
-
-# Level 0 keeps napari-friendly chunks even when tiles are much larger; the
-# cap has to divide the tile shape so merge workers still write whole chunks.
-out_chunks = capped_output_chunks(staged.chunks, (16, 1024, 1024))
+if in_place:
+    # The tiles already sit in labels/<name>/0, so the merge rewrites them
+    # where they are: no scratch store, and one full write of the volume less.
+    # Safe because the boundary scan finishes before any chunk is rewritten.
+    out_chunks = None
+else:
+    # Level 0 keeps napari-friendly chunks even when tiles are much larger;
+    # the cap must divide the tile so workers still write whole chunks.
+    root = zarr.open_group(image_store, mode="a")
+    parent = root.require_group("labels")
+    if label_name in parent:
+        del parent[label_name]
+    parent.require_group(label_name)
+    out_chunks = capped_output_chunks(staged.chunks, (16, 1024, 1024))
 
 _, n_objects = merge_tile_labels(
-    stage_path(work_dir, label_name),
-    write_to=label_group,
-    input_component="staged",
-    output_component="0",
+    target_path,
+    write_to=label_group if not in_place else target_path,
+    input_component=target_component,
+    output_component="0" if not in_place else target_component,
     output_chunks=out_chunks,
     sequential_labels=cfg.get("sequential_labels", True),
     n_workers=cfg.get("merge_workers") or default_workers,
@@ -88,11 +96,12 @@
     n_objects=n_objects,
 )
 
-shutil.rmtree(stage_path(work_dir, label_name), ignore_errors=True)
-# Also drop the checkpoint's completion sentinel (stage.zarr.done): the
-# "prepare" rule's stage=touch(STAGE_OK) output must not outlive the store it
-# claims exists, or a future rerun (e.g. re-segmenting for new labels) skips
-# "prepare" and "segment" tries to open a stage.zarr that's already gone.
+if not in_place:
+    shutil.rmtree(stage_path(work_dir, label_name), ignore_errors=True)
+# Drop the checkpoint's completion sentinel (stage.zarr.done): the "prepare"
+# rule's stage=touch(STAGE_OK) output must not outlive what it claims exists,
+# or a future rerun (e.g. re-segmenting for new labels) skips "prepare" and
+# "segment" writes into a target that is already gone or already merged.
 Path(f"{stage_path(work_dir, label_name)}.done").unlink(missing_ok=True)
 print(f"[patchworks] labels written to {group}")
 open(snakemake.output[0], "w").close()  # noqa: F821

workflow/scripts/prepare_tiles.py

@@ -5,13 +5,15 @@
 from pathlib import Path
 
 import numpy as np
+import zarr
 
 from patchworks import (
     auto_empty_threshold,
     auto_tile_shape,
     auto_tile_shape_cellpose,
     block_for_tile,
     build_occupancy_map,
+    capped_output_chunks,
     create_stage,
     normalize_overlap,
     spatial_tiles,
@@ -20,6 +22,9 @@
 
 from _pw import open_image, stage_path, start_log, validate_config
 
+# Chunking ceiling for the written labels, so a viewer can page them lazily.
+LABEL_CHUNK_CAP = (16, 1024, 1024)
+
 start_log(snakemake.log[0])  # noqa: F821
 cfg = snakemake.config  # noqa: F821
 work_dir = cfg["work_dir"]
@@ -134,7 +139,31 @@
     for i in range(0, len(occupied), tiles_per_job)
 ]
 
-create_stage(stage_path(work_dir, label_name), image.shape, tile_shape)
+# Where the segment jobs write. When the tile is already within the label
+# chunk cap, they can write straight into the label group's level 0 and the
+# merge relabels it in place -- one full write of the volume less than
+# staging to a scratch store and copying it across. An oversized tile keeps
+# the scratch store, so level 0 can still be chunked for lazy viewing.
+image_store = str(Path(work_dir) / "image.zarr")
+in_place = capped_output_chunks(tile_shape, LABEL_CHUNK_CAP) == tuple(
+    tile_shape
+)
+if in_place:
+    target_path, target_component = f"{image_store}/labels/{label_name}", "0"
+    root = zarr.open_group(image_store, mode="a")
+    parent = root.require_group("labels")
+    if label_name in parent:
+        del parent[label_name]
+    parent.require_group(label_name)
+    print("[patchworks] segmenting straight into the label group (in place)")
+else:
+    target_path, target_component = stage_path(work_dir, label_name), "staged"
+    print(
+        f"[patchworks] tile {tuple(tile_shape)} exceeds the label chunk cap "
+        f"{LABEL_CHUNK_CAP}; staging to a scratch store so level 0 stays "
+        "chunked for viewing"
+    )
+create_stage(target_path, image.shape, tile_shape, component=target_component)
 
 Path(work_dir, label_name, "tiles.json").write_text(
     json.dumps(
@@ -145,6 +174,9 @@
             "occupied": occupied,
             "tiles_per_job": tiles_per_job,
             "batches": batches,
+            "target_path": target_path,
+            "target_component": target_component,
+            "in_place": in_place,
         },
         indent=2,
     )

workflow/scripts/segment_tile.py

@@ -11,7 +11,7 @@
 
 from patchworks import stage_tile
 
-from _pw import build_fn, load_tiles_json, open_image, stage_path, start_log
+from _pw import build_fn, load_tiles_json, open_image, start_log
 
 start_log(snakemake.log[0])  # noqa: F821
 cfg = snakemake.config  # noqa: F821
@@ -25,7 +25,11 @@
 
 # Built once for the whole batch: this is what makes the model load amortize.
 fn = build_fn(cfg)
-stage = stage_path(work_dir, label_name)
+# prepare decides where tiles land: the label group's level 0 directly when
+# the tile fits the chunk cap (the merge then relabels it in place), else a
+# scratch stage store.
+stage = manifest["target_path"]
+component = manifest.get("target_component", "staged")
 tile_shape = tuple(manifest["tile_shape"])
 
 counts = {}
@@ -38,6 +42,7 @@
         tile_shape=tile_shape,
         # Scalar (older manifests) or per-axis list; stage_tile normalizes both.
         overlap=manifest["overlap"],
+        component=component,
     )
     print(f"[patchworks] segmented tile {index}: {counts[index]} label(s)")