perf(merge): ⚡ parallelize the boundary scan and give every pool its cores

Three places were leaving work on the table.

_scan_touching_pairs was serial. Each task reads its own boundary slab and
produces its own pairs, so it was embarrassingly parallel already; it now
runs on a thread pool sized like the rest of the merge. (Planned with the
offsets work, missed at the time.)

The streaming pyramid called _stream_strided_level without n_workers, so it
took the signature default of 4 threads no matter what the job was granted --
8x short on a 32-core merge.

tile_process, the direct API, never got the merge rework: it still ran
zarr_native_merge and *then* relabel_sequential_zarr, paying an extra full
read+write of the volume plus a Python set of every label id. It now passes
sequential= so the renumbering folds into the merge's own LUT, exactly as the
workflow does. The library path was slower than the workflow path.

Also corrects the _level_chunks docstring: the floor overrides the
one-source-chunk-per-task property at deep levels, where an output chunk
covers a few source chunks instead of one. Bounded by a small constant either
way, which is the property that matters.

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

Author: lguerard

src/patchworks/_core.py

@@ -15,7 +15,6 @@
 from ._cluster import _client_is_in_process, _distributed_client
 from ._io import auto_empty_threshold, load_ome_zarr
 from ._merge import zarr_native_merge
-from ._relabel import relabel_sequential_zarr
 
 logger = logging.getLogger(__name__)
 
@@ -606,17 +605,19 @@ def _cleanup_stage():
             tempfile.mkdtemp(prefix="bb_merge_"), "merged.zarr"
         )
 
+    # sequential=True folds the contiguous renumbering into the merge's own
+    # LUT, so it costs a np.unique over the object count rather than the extra
+    # full read+write (plus a Python set of every id) that a separate
+    # relabel_sequential_zarr pass would.
     zarr_native_merge(
         stage_path,
         "staged",
         _merge_out,
         output_component,
         n_workers=_nw,
         show_progress=progress,
+        sequential=sequential_labels,
     )
-    if sequential_labels:
-        logger.info("Relabelling to contiguous ids…")
-        relabel_sequential_zarr(_merge_out, output_component)
     _cleanup_stage()
 
     merged = da.from_zarr(_merge_out, component=output_component)

src/patchworks/_merge.py

@@ -22,6 +22,7 @@
 import logging
 import os
 import tempfile
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import nullcontext as _nullcontext
 from itertools import product as _iproduct
 from multiprocessing import Pool as _Pool
@@ -155,6 +156,7 @@ def _scan_touching_pairs(
     component: str,
     chunk_shape: tuple[int, ...],
     label_offsets: "np.ndarray | None" = None,
+    n_workers: int = 1,
 ) -> np.ndarray:
     """Scan chunk-boundary slabs; return (N, 2) int64 array of touching pairs.
 
@@ -178,6 +180,9 @@ def _scan_touching_pairs(
         come out global without the store ever being rewritten. Each read is
         confined to one chunk column, so both sides are a single chunk and the
         offsets are plain scalars.
+    n_workers : int
+        Threads to scan with. Each task reads its own slab, so this is
+        embarrassingly parallel and was needlessly serial.
 
     Returns
     -------
@@ -189,36 +194,52 @@ def _scan_touching_pairs(
     arr = root[component]
     shape = arr.shape
     n_per_dim = [(s + c - 1) // c for s, c in zip(shape, chunk_shape)]
-    specs = _boundary_face_specs(shape, chunk_shape)
-    all_pairs: list[np.ndarray] = []
-    for ax, pos in specs:
-        # tile the face dimensions using chunk_shape columns
+
+    # One task per chunk-column of one boundary face. Each reads its own slab
+    # and produces its own pairs, so they are independent.
+    tasks: list[tuple[int, int, tuple[int, ...]]] = []
+    for ax, pos in _boundary_face_specs(shape, chunk_shape):
         face_axes = [a for a in range(arr.ndim) if a != ax]
         face_ranges = [range(0, shape[a], chunk_shape[a]) for a in face_axes]
         for offsets in _iproduct(*face_ranges):
-            sl: list = [slice(None)] * arr.ndim
-            sl[ax] = slice(pos - 1, pos + 1)
+            tasks.append((ax, pos, offsets))
+
+    def _one(task: tuple[int, int, tuple[int, ...]]) -> "np.ndarray | None":
+        ax, pos, offsets = task
+        face_axes = [a for a in range(arr.ndim) if a != ax]
+        sl: list = [slice(None)] * arr.ndim
+        sl[ax] = slice(pos - 1, pos + 1)
+        for a, off in zip(face_axes, offsets):
+            sl[a] = slice(off, min(off + chunk_shape[a], shape[a]))
+        slab = np.moveaxis(np.asarray(arr[tuple(sl)]), ax, 0)
+        a_vals = slab[0].ravel().astype(np.int64)
+        b_vals = slab[1].ravel().astype(np.int64)
+        if label_offsets is not None:
+            grid = [0] * arr.ndim
             for a, off in zip(face_axes, offsets):
-                sl[a] = slice(off, min(off + chunk_shape[a], shape[a]))
-            slab = np.moveaxis(np.asarray(arr[tuple(sl)]), ax, 0)
-            a_vals = slab[0].ravel().astype(np.int64)
-            b_vals = slab[1].ravel().astype(np.int64)
-            if label_offsets is not None:
-                grid = [0] * arr.ndim
-                for a, off in zip(face_axes, offsets):
-                    grid[a] = off // chunk_shape[a]
-                grid[ax] = pos // chunk_shape[ax]
-                b_idx = int(np.ravel_multi_index(tuple(grid), n_per_dim))
-                grid[ax] -= 1  # the chunk on the near side of the boundary
-                a_idx = int(np.ravel_multi_index(tuple(grid), n_per_dim))
-                a_vals[a_vals > 0] += label_offsets[a_idx]
-                b_vals[b_vals > 0] += label_offsets[b_idx]
-            mask = (a_vals > 0) & (b_vals > 0) & (a_vals != b_vals)
-            if mask.any():
-                pairs = np.sort(
-                    np.stack([a_vals[mask], b_vals[mask]], axis=1), axis=1
-                )
-                all_pairs.append(np.unique(pairs, axis=0))
+                grid[a] = off // chunk_shape[a]
+            grid[ax] = pos // chunk_shape[ax]
+            b_idx = int(np.ravel_multi_index(tuple(grid), n_per_dim))
+            grid[ax] -= 1  # the chunk on the near side of the boundary
+            a_idx = int(np.ravel_multi_index(tuple(grid), n_per_dim))
+            a_vals[a_vals > 0] += label_offsets[a_idx]
+            b_vals[b_vals > 0] += label_offsets[b_idx]
+        mask = (a_vals > 0) & (b_vals > 0) & (a_vals != b_vals)
+        if not mask.any():
+            return None
+        pairs = np.sort(np.stack([a_vals[mask], b_vals[mask]], axis=1), axis=1)
+        return np.unique(pairs, axis=0)
+
+    nw = max(1, min(n_workers, len(tasks)))
+    if nw <= 1:
+        results = [_one(t) for t in tasks]
+    else:
+        # Reads and decompression release the GIL, so threads scale here and
+        # nothing has to be pickled across processes.
+        with ThreadPoolExecutor(max_workers=nw) as pool:
+            results = list(pool.map(_one, tasks))
+
+    all_pairs = [r for r in results if r is not None]
     if not all_pairs:
         return np.empty((0, 2), dtype=np.int64)
     return np.unique(np.vstack(all_pairs), axis=0)
@@ -545,7 +566,11 @@ def zarr_native_merge(
     n_faces = len(_boundary_face_specs(shape, chunk_shape))
     logger.info("zarr_native_merge: scanning %d boundary faces…", n_faces)
     pairs = _scan_touching_pairs(
-        staged_path, staged_component, chunk_shape, label_offsets=offsets
+        staged_path,
+        staged_component,
+        chunk_shape,
+        label_offsets=offsets,
+        n_workers=n_workers,
     )
     logger.info(
         "zarr_native_merge: %d touching pairs → building LUT", len(pairs)

src/patchworks/plugins/ome_zarr.py

@@ -56,6 +56,7 @@
 import numpy as np
 import zarr
 
+from .._chunks import cpu_allocation
 from .._io import load_ome_zarr
 
 logger = logging.getLogger(__name__)
@@ -290,12 +291,16 @@ def _level_chunks(
     axes: str,
     level_shape: tuple[int, ...],
 ) -> tuple[int, ...]:
-    """Chunking for a strided level that keeps one source chunk per task.
+    """Chunking for a strided level that keeps the per-task read bounded.
 
     Choosing ``ceil(src_chunk / stride)`` makes each output chunk the exact
     image of one source chunk, so a task reads one chunk and writes one chunk
-    with nothing shared between tasks. Capped by :data:`_CHUNK_CAP` and floored
-    so deep levels don't fragment.
+    with nothing shared between tasks.
+
+    The floor can override that at deep levels, where ``ceil(src / stride)``
+    falls below it — an output chunk then covers a few source chunks instead
+    of one. Still bounded by a small constant, which is the property that
+    matters; those levels are a fraction of the volume anyway.
     """
     out = []
     for c, st, a, s in zip(src_chunks, strides, axes, level_shape):
@@ -621,7 +626,9 @@ def _write_pyramid(
                         )
                     ),
                 )
-            _stream_strided_level(src_arr, dst_arr, strides)
+            _stream_strided_level(
+                src_arr, dst_arr, strides, n_workers=cpu_allocation()
+            )
         scale = [base_scale[k] * (strides[k] ** i) for k in range(len(axes))]
         datasets.append(_dataset(str(i), scale))
         logger.info("pyramid level %d: shape=%s", i, next_shape)