fix(core): make per-tile labels globally unique before merging

src/patchworks/_distributed.py

@@ -68,7 +68,8 @@ def create_stage(
     component : str, optional
         Array name inside the store (default ``"staged"``).
     dtype : data-type, optional
-        Label dtype (default ``int32``).
+        Label dtype (default ``int32``). Tiles write local labels; the merge's
+        first pass renumbers them to a compact global range that fits int32.
 
     Returns
     -------
@@ -135,6 +136,8 @@ def stage_tile(
         slice(left, out.shape[i] - right)
         for i, (left, right) in enumerate(trims)
     )
+    # Local labels (1..N) are fine here — they collide across tiles, but the
+    # merge's first pass makes them globally unique before stitching.
     dst = zarr.open_group(str(stage_path), mode="r+")[component]
     dst[sl] = out[sel].astype(dst.dtype)
     return index

src/patchworks/_merge.py

@@ -103,15 +103,15 @@ def _relabel_chunk_worker(chunk_slice: tuple) -> None:
     block = np.asarray(src[chunk_slice], dtype=np.int64)
     max_b = int(block.max())
     if max_b == 0:
-        dst[chunk_slice] = block.astype(np.int32)
+        dst[chunk_slice] = block.astype(dst.dtype)
         return
     lut = _merge_lut
     if max_b < len(lut):
         out = lut[block]
     else:
         ext = np.arange(len(lut), max_b + 1, dtype=np.int64)
         out = np.concatenate([lut, ext])[block]
-    dst[chunk_slice] = out.astype(np.int32)
+    dst[chunk_slice] = out.astype(dst.dtype)
 
 
 def _boundary_face_specs(
@@ -238,9 +238,9 @@ def _build_relabel_lut(pairs: np.ndarray, max_label: int) -> np.ndarray:
 
 
 def _create_zarr_label_array(
-    group: zarr.Group, name: str, shape: tuple, chunks: tuple
+    group: zarr.Group, name: str, shape: tuple, chunks: tuple, dtype=np.int32
 ) -> zarr.Array:
-    """Create (replacing any existing) an int32 label array in *group*.
+    """Create (replacing any existing) a label array in *group*.
 
     Parameters
     ----------
@@ -252,6 +252,9 @@ def _create_zarr_label_array(
         Array shape.
     chunks : tuple
         Chunk shape.
+    dtype : data-type, optional
+        Label dtype (default ``int32``). Matched to the staged store so the
+        merge keeps wide (block-encoded) ids intact before they are compacted.
 
     Returns
     -------
@@ -261,14 +264,56 @@ def _create_zarr_label_array(
     if name in group:
         del group[name]
     if _ZARR_V3:
-        return group.create_array(
-            name, shape=shape, chunks=chunks, dtype=np.int32
-        )
+        return group.create_array(name, shape=shape, chunks=chunks, dtype=dtype)
     return group.zeros(
-        name, shape=shape, chunks=chunks, dtype=np.int32, overwrite=True
+        name, shape=shape, chunks=chunks, dtype=dtype, overwrite=True
     )
 
 
+def _make_globally_unique(arr, shape: tuple, chunk_shape: tuple) -> int:
+    """Renumber per-chunk local labels to a globally unique, compact range.
+
+    Each tile writes local labels (``1..N``) that repeat across tiles, so the
+    same value means different objects in different tiles. This streams the
+    chunks in row-major order and remaps every chunk's non-zero labels to a
+    fresh contiguous block ``[base+1, …]``, leaving the store globally unique
+    with ``max_label == total objects`` (compact, so the relabel LUT stays
+    ``O(n_objects)``). Background (0) stays 0. In place, one chunk in RAM.
+
+    Parameters
+    ----------
+    arr : zarr.Array
+        Writable staged label array.
+    shape : tuple
+        Array shape.
+    chunk_shape : tuple
+        Chunk (= tile) shape.
+
+    Returns
+    -------
+    int
+        New maximum label (number of objects across all tiles, before the
+        cross-boundary merge fuses touching pairs).
+    """
+    n_per_dim = [(s + c - 1) // c for s, c in zip(shape, chunk_shape)]
+    base = 0
+    for idx in _iproduct(*[range(n) for n in n_per_dim]):
+        sl = tuple(
+            slice(i * c, min((i + 1) * c, s))
+            for i, c, s in zip(idx, chunk_shape, shape)
+        )
+        block = np.asarray(arr[sl])
+        uniq = np.unique(block)
+        uniq = uniq[uniq > 0]
+        if uniq.size == 0:
+            continue
+        lut = np.zeros(int(uniq[-1]) + 1, dtype=np.int64)
+        lut[uniq] = np.arange(1, uniq.size + 1) + base
+        arr[sl] = lut[block].astype(arr.dtype)
+        base += int(uniq.size)
+    return base
+
+
 def zarr_native_merge(
     staged_path: str,
     staged_component: str,
@@ -302,13 +347,17 @@ def zarr_native_merge(
     -------
     None
     """
-    root = zarr.open_group(staged_path, mode="r")
+    root = zarr.open_group(staged_path, mode="r+")
     arr = root[staged_component]
     shape, chunk_shape = arr.shape, arr.chunks
 
-    max_label = int(
-        da.from_zarr(staged_path, component=staged_component).max().compute()
-    )
+    # Pass 0 — make labels globally unique. Each tile writes local labels
+    # (1..N) that collide across tiles (every tile has a "1"); without this the
+    # boundary merge would fuse unrelated objects that merely share a value.
+    # Stream chunks in order, remapping each chunk's labels into a fresh
+    # contiguous range [base+1, …] — unique *and* compact (so the relabel LUT
+    # stays O(n_objects), not O(n_voxels)).
+    max_label = _make_globally_unique(arr, shape, chunk_shape)
     logger.info(
         "zarr_native_merge: shape=%s chunks=%s max_label=%d",
         shape,
@@ -330,7 +379,11 @@ def zarr_native_merge(
     )
 
     out_root = zarr.open_group(out_path, mode="a")
-    _create_zarr_label_array(out_root, out_component, shape, chunk_shape)
+    # Match the staged dtype so block-encoded (wide) ids survive the merge;
+    # they are compacted to a small range afterwards by relabel_sequential_zarr.
+    _create_zarr_label_array(
+        out_root, out_component, shape, chunk_shape, dtype=arr.dtype
+    )
 
     n_per_dim = [(s + c - 1) // c for s, c in zip(shape, chunk_shape)]
     chunk_slices = [

tests/test_distributed.py

@@ -43,3 +43,33 @@ def test_stage_then_merge_stitches_boundary(tmp_path):
     ).compute()
     ids = np.unique(merged[merged > 0])
     assert ids.size == 1, f"object split into {ids.size} labels"
+
+
+def test_separate_objects_keep_distinct_labels(tmp_path):
+    """Different objects in different tiles must NOT collapse to one label.
+
+    Each tile produces local labels (1..N); without the merge's global-uniqueness
+    pass every tile's "1" would fuse into a single object.
+    """
+    img = np.zeros((16, 32), "uint16")
+    # four separate objects, one per (16x16) tile, none touching a boundary
+    img[3:6, 3:6] = 500
+    img[3:6, 19:22] = 500
+    img[10:13, 3:6] = 500
+    img[10:13, 19:22] = 500
+
+    stage = str(tmp_path / "stage.zarr")
+    tile = (16, 16)
+    create_stage(stage, img.shape, tile)
+    for i in range(len(spatial_tiles(img.shape, tile))):
+        stage_tile(img, _fn, stage, i, tile_shape=tile, overlap=4)
+
+    merged = merge_tile_labels(
+        stage,
+        write_to=str(tmp_path / "out.zarr"),
+        input_component="staged",
+        sequential_labels=True,
+    ).compute()
+    ids = np.unique(merged[merged > 0])
+    assert ids.size == 4, f"expected 4 distinct objects, got {ids.size}"
+    assert set(ids.tolist()) == {1, 2, 3, 4}  # contiguous after relabel