fix(ome_zarr): 🐛 guard in-place re-merges and match NGFF layout to zarr version

Four fixes, three of them found by comparing against napari-tmidas.

An in-place merge destroys its own input, so running it twice would add each
tile's offset to ids that are already global and could land two unrelated
objects on the same id. The merge now records its progress on the array:
"running" (a previous attempt died mid-relabel, so the data is part local and
part global) refuses and asks for a re-segmentation; "done" is a no-op, so a
failure *after* the relabel can still be retried cheaply. An earlier retry
test passed only by luck.

The NGFF metadata layout was wrong for the store being written. NGFF 0.4 is
defined over zarr v2 with its keys at the top level; 0.5 is the v3 revision
and nests them under "ome" with the version there. patchworks wrote v3 data
with 0.4's top-level layout, which matches neither -- a strict 0.5 reader
looking under attributes.ome found nothing. The layout now follows the zarr
version, and every reader goes through read_ngff_attr, which accepts both, so
stores written by earlier versions still load.

_default_axes guesses from the number of dimensions, which is unambiguous up
to 3-D but silently labels a (t, z, y, x) array "czyx". It now warns. Files
never reach this path -- their axes come from the reader's dimension
metadata, which is what napari-tmidas does and what patchworks already did.

The compressor is pinned rather than inherited from zarr's default (zstd
either way today, but the stores should not change silently if that moves).

Verified end to end on both merge routes, through the napari plugin's label
discovery, and against napari-chunked-regionprops (17/17, measurements
matching direct voxel counts).

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

Author: lguerard

docs/api/plugins/ome_zarr.md

@@ -24,3 +24,15 @@ resolution, matching anisotropic microscopy stacks.
 ## register_labels
 
 ::: patchworks.plugins.ome_zarr.register_labels
+
+## NGFF metadata layout
+
+NGFF 0.4 is defined over zarr v2 and puts its keys at the top level; 0.5 is
+the zarr-v3 revision and nests them under `ome`. patchworks writes whichever
+matches the store, and reads both.
+
+::: patchworks.plugins.ome_zarr.ngff_version
+
+::: patchworks.plugins.ome_zarr.read_ngff_attr
+
+::: patchworks.plugins.ome_zarr.write_ngff_attrs

docs/guide/merging.md

@@ -94,9 +94,19 @@ image that is most of the volume.
 
 When the merge's output *is* its input, this pass rewrites the array in
 place. That is safe because the boundary scan (step 3) has already finished,
-so nothing still needs the original ids. The trade-off is restartability: a
-job killed part-way through leaves a half-relabelled array, where a separate
-output store would have left the input intact to redo from.
+so nothing still needs the original ids.
+
+Because an in-place merge destroys its own input, it records how far it got
+on the array itself, and refuses to guess on a re-run:
+
+| State found | What happens |
+| --- | --- |
+| nothing recorded | fresh tile-local ids — merge normally |
+| `running` | a previous attempt died mid-relabel, so the array is part local and part global. **Refuses**: re-segment to rebuild it. |
+| `done` | already merged — a no-op, so a failure *after* the relabel (the pyramid, say) can simply be retried |
+
+Without that, a second pass would add each tile's offset to ids that are
+already global, which can land two unrelated objects on the same id.
 
 ## Using the merge step standalone
 

docs/guide/ome_zarr_napari.md

@@ -22,9 +22,23 @@ natively. Pass `write_to="other.zarr"` to instead write a separate
 single-resolution label store, or `output_component="cells"` to name the label
 image.
 
-The label pyramid is built lazily (`da.to_zarr`, streamed chunk by chunk), so
-it stays OOM-safe even for terabyte volumes. Control it with `pyramid_levels`
-and `pyramid_downscale`.
+The label pyramid is streamed chunk by chunk, so it stays OOM-safe even for
+terabyte volumes. Control it with `pyramid_levels` and `pyramid_downscale`.
+
+!!! note "NGFF version"
+    The metadata layout follows the zarr format being written: **NGFF 0.5**
+    (keys nested under an `ome` attribute) on zarr v3, **0.4** (keys at the
+    top level) on v2. Writing v3 data with 0.4's top-level layout — which
+    earlier versions did — matches neither revision, and a strict 0.5 reader
+    finds nothing there.
+
+    Reading accepts **both** layouts, so stores written by any patchworks
+    version still load. Use `read_ngff_attr()` if you need to inspect the
+    metadata yourself rather than indexing `attrs["multiscales"]` directly.
+
+    patchworks' own hints (`n_objects`, `sequential_labels`) are not NGFF
+    keys, so they stay at the top level where a consumer finds them without
+    knowing the layout.
 
 ## Why a pyramid?
 
@@ -107,6 +121,21 @@ input — see [Cluster usage](snakemake.md).
     to_ome_zarr("scan.ims", "scan.zarr", reuse_pyramid=True)
     ```
 
+### Axis names
+
+For a file, axes come from the reader's own dimension metadata (bioio's
+`dims.order`, Imaris, or the `sequence_pattern` groups) — not guessed. Only
+non-spatial singleton axes are dropped, so a genuine `z=1` survives.
+
+For a **bare array** there is nothing to read, so axes are inferred from the
+number of dimensions (`yx`, `zyx`, `czyx`, `tczyx`). That is unambiguous up to
+3-D; from 4-D the leading axis could be channel or time, and patchworks logs a
+warning saying what it guessed. Pass `axes=` to settle it:
+
+```python
+to_ome_zarr(arr, "out.zarr", axes="tzyx")   # not the czyx it would assume
+```
+
 ### Pixel calibration
 
 The physical voxel size is read from the input — bioio's `physical_pixel_sizes`,

src/patchworks/_distributed.py

@@ -18,6 +18,7 @@
 import numpy as np
 import zarr
 
+from ._io import zarr_compressor_kwargs
 from ._relabel import relabel_sequential_array
 
 Overlap = Union[int, Sequence[int]]
@@ -116,7 +117,11 @@ def create_stage(
     """
     root = zarr.open_group(str(stage_path), mode="w")
     root.create_array(
-        name=component, shape=shape, chunks=tile_shape, dtype=dtype
+        name=component,
+        shape=shape,
+        chunks=tile_shape,
+        dtype=dtype,
+        **zarr_compressor_kwargs(),
     )
     return str(stage_path)
 

src/patchworks/_io.py

@@ -15,6 +15,35 @@
 _ZARR_V3 = int(zarr.__version__.split(".")[0]) >= 3
 
 
+_ZARR_V3 = int(zarr.__version__.split(".")[0]) >= 3
+
+
+def zarr_compressor_kwargs() -> dict:
+    """Keyword arguments pinning the compression codec for a new array.
+
+    zstd is already zarr v3's default, but relying on a library default means
+    the stores patchworks writes change silently if that default ever moves.
+    Labels in particular are highly compressible, so this is worth stating.
+
+    Returns
+    -------
+    dict
+        ``compressors=``/``compressor=`` as the installed zarr expects, or
+        empty if the codec cannot be built (then the default applies).
+    """
+    try:
+        if _ZARR_V3:
+            from zarr.codecs import ZstdCodec
+
+            return {"compressors": (ZstdCodec(level=1),)}
+        import numcodecs
+
+        return {"compressor": numcodecs.Zstd(level=1)}
+    except Exception:  # pragma: no cover - depends on the installed zarr
+        logger.debug("could not pin a compressor; using zarr's default")
+        return {}
+
+
 def load_ome_zarr(
     store_path: Union[str, Path],
     channel: int | None = 0,

src/patchworks/_merge.py

@@ -34,6 +34,7 @@
 import zarr
 
 from ._chunks import cpu_allocation
+from ._io import zarr_compressor_kwargs
 
 try:
     from tqdm.auto import tqdm as _tqdm
@@ -45,6 +46,12 @@
 _ZARR_V3 = int(zarr.__version__.split(".")[0]) >= 3
 _LUT_WARN_THRESHOLD = 100_000_000  # warn when max_label > 100 M (LUT > 800 MB)
 
+# Attributes recording how far an in-place merge got. Needed because an
+# in-place merge destroys its own input: a second pass over already-global ids
+# would add the per-tile offsets again and can collide unrelated objects.
+_MERGE_STATE = "patchworks_merge_state"
+_MERGE_COUNT = "patchworks_n_objects"
+
 # Per-worker globals set by _init_worker.
 # LUT is memory-mapped from disk so it is shared read-only across all workers
 # (OS page cache, no per-process copy). Passing the LUT directly via pickle
@@ -335,10 +342,13 @@ def _create_zarr_label_array(
     """
     if name in group:
         del group[name]
+    kwargs = zarr_compressor_kwargs()
     if _ZARR_V3:
-        return group.create_array(name, shape=shape, chunks=chunks, dtype=dtype)
+        return group.create_array(
+            name, shape=shape, chunks=chunks, dtype=dtype, **kwargs
+        )
     return group.zeros(
-        name, shape=shape, chunks=chunks, dtype=dtype, overwrite=True
+        name, shape=shape, chunks=chunks, dtype=dtype, overwrite=True, **kwargs
     )
 
 
@@ -635,6 +645,29 @@ def zarr_native_merge(
                 "output_chunks cannot differ from the source chunking when "
                 "merging in place (the array is rewritten, not recreated)"
             )
+        # Running this twice over the same array would add the per-tile
+        # offsets to ids that are already global, which can land two unrelated
+        # objects on the same id. The source is destroyed as we go, so the
+        # state has to be recorded rather than inferred.
+        state = arr.attrs.get(_MERGE_STATE)
+        if state == "running":
+            raise RuntimeError(
+                f"{out_path}/{out_component} was left half-merged by an "
+                "earlier attempt, so its labels are part local and part "
+                "global and cannot be merged again. Re-run the segmentation "
+                "for this label (delete the label group and its tiles.json) "
+                "to rebuild it from scratch."
+            )
+        if state == "done":
+            n = arr.attrs.get(_MERGE_COUNT)
+            logger.info(
+                "zarr_native_merge: %s/%s is already merged (%s objects); "
+                "nothing to do",
+                out_path,
+                out_component,
+                n,
+            )
+            return n
         logger.info("zarr_native_merge: relabeling in place, no second store")
 
     out_root = zarr.open_group(out_path, mode="a")
@@ -699,6 +732,12 @@ def zarr_native_merge(
     np.save(lut_path, lut)
     del lut  # parent no longer needs it; workers load via mmap
 
+    if in_place:
+        # From here the source is being overwritten. Record that, so a crash
+        # leaves evidence the array is half-converted rather than looking like
+        # untouched input to the next attempt.
+        arr.attrs[_MERGE_STATE] = "running"
+
     try:
         if n_w <= 1:
             _init_worker(
@@ -731,6 +770,13 @@ def zarr_native_merge(
 
         shutil.rmtree(_lut_dir, ignore_errors=True)
 
+    if in_place:
+        # Every chunk is now global. Recording the count as well makes a retry
+        # after a *later* failure (the pyramid, say) a cheap no-op instead of
+        # forcing a full re-segmentation.
+        arr.attrs[_MERGE_COUNT] = n_objects
+        arr.attrs[_MERGE_STATE] = "done"
+
     return n_objects
 
 

src/patchworks/plugins/napari.py

@@ -32,6 +32,7 @@
 import zarr
 
 from .._io import load_ome_zarr
+from .ome_zarr import read_ngff_attr
 
 logger = logging.getLogger(__name__)
 
@@ -84,7 +85,7 @@ def _has_multiscales(path: Union[str, Path]) -> bool:
         True if the group has a ``multiscales`` attribute.
     """
     root = zarr.open_group(str(path), mode="r")
-    return "multiscales" in root.attrs
+    return read_ngff_attr(root.attrs, "multiscales") is not None
 
 
 def _multiscale_levels(
@@ -105,7 +106,7 @@ def _multiscale_levels(
         One lazy array per resolution level.
     """
     root = zarr.open_group(str(path), mode="r")
-    datasets = root.attrs["multiscales"][0]["datasets"]
+    datasets = read_ngff_attr(root.attrs, "multiscales")[0]["datasets"]
     return [
         load_ome_zarr(path, channel=channel, level=i)
         for i in range(len(datasets))
@@ -231,7 +232,7 @@ def _inner_label_names(store: Union[str, Path]) -> list[str]:
         grp = zarr.open_group(f"{store}/labels", mode="r")
     except Exception:
         return []
-    return list(grp.attrs.get("labels", []))
+    return list(read_ngff_attr(grp.attrs, "labels", []) or [])
 
 
 def _resolve_labels(