[DO NOT MERGE] Parallel map

CHANGELOG.md

@@ -1,5 +1,30 @@
 # Changelog
 
+## [Unreleased]
+
+### API Breaking Changes
+
+- `MapperProtocol` and `BasicMapper` are re-typed around the new `IterUnit`: a mapper now receives `(func, units)` instead of `(func, getters, setters)` and returns the results as a list ordered by `unit.index`. A unit whose `setter` is `None` is a read-only unit (compute and collect, never write), no longer an error. `map_as_numpy`/`map_as_dask` call sites are unaffected — only custom mapper implementations need updating:
+
+```python
+def __call__(self, func, units):                       # was (self, func, getters, setters)
+    return [func(unit.getter()) for unit in units]     # was: zip getters/setters, return None
+```
+
+### Behaviour changes
+
+- `by_chunks()` gained a `grid="write" | "read"` parameter, defaulting to `"write"`: tiles now align to the output image's write granularity (shard shape when sharded, chunk shape otherwise) instead of the input chunk grid, so the resulting ROIs pass `check_if_chunks_overlap`. Identical for default-derived outputs; different when the output has custom `chunks=`/`shards=`. Read-only iterators fall back to the input grid; `grid="read"` restores the old tiling.
+
+### Features
+
+- `reduce_as_numpy`/`reduce_as_dask` on all iterators: apply a function to every ROI and collect the results (any return type, e.g. a `DataFrame` per ROI) without writing anything — `results[i]` corresponds to `iterator.rois[i]` regardless of mapper execution order.
+- `IterUnit`: the unit of work a mapper schedules, carrying `index`, `roi`, `getter`, `setter` and `write_footprint` (the chunk/shard rectangle the unit writes — the conflict currency for parallel mappers).
+
+### Fix
+
+- `check_if_chunks_overlap`/`require_no_chunks_overlap` measured the *read* side: slicing from the getters against the input image's chunk grid. They now measure the write target — slicing from the setters against the output array's write granularity, which is the shard shape when the output is sharded (writes are atomic per shard object). This is a behaviour change: layouts that previously passed the gate (finely chunked input with a coarser-chunked or sharded output) may now correctly raise, and read-only iterators now always report no overlap.
+- Dask writes (`set_array(..., mode="dask")`, masked dask setters, pyramid consolidation) could silently drop data when the patch's block grid was misaligned with the target's chunk grid — or, for sharded targets, its shard grid: `da.store(lock=False)` let two blocks read-modify-write the same chunk concurrently. All dask store flushes now serialise behind a shared lock; block compute stays parallel.
+
 ## [v1.0.0]
 
 First stable release. Everything deprecated in `v0.5.0` (each warned "will be removed in `ngio=0.6`") is now removed — that release became `1.0.0`.

src/ngio/__init__.py

@@ -38,6 +38,7 @@
     BasicMapper,
     FeatureExtractorIterator,
     ImageProcessingIterator,
+    IterUnit,
     MapperProtocol,
     MaskedSegmentationIterator,
     SegmentationIterator,
@@ -82,6 +83,7 @@
     "Image",
     "ImageInWellPath",
     "ImageProcessingIterator",
+    "IterUnit",
     "Label",
     "MapperProtocol",
     "MaskedImage",

src/ngio/common/_locks.py

@@ -0,0 +1,6 @@
+from dask.utils import SerializableLock
+
+# Serialises zarr chunk read-modify-writes issued by da.store across all ngio
+# dask write sites; block compute stays parallel. SerializableLock rather than
+# threading.Lock so graphs survive pickling under a distributed client.
+DASK_STORE_LOCK = SerializableLock()

src/ngio/common/_pyramid.py

@@ -8,6 +8,7 @@
 import zarr
 from pydantic import BaseModel, ConfigDict, model_validator
 
+from ngio.common._locks import DASK_STORE_LOCK
 from ngio.common._zoom import (
     InterpolationOrder,
     _zoom_inputs_check,
@@ -45,7 +46,10 @@ def _on_disk_dask_zoom(
     # re-derives chunks via normalize_chunks(chunks="auto", ...) and warns
     # (treated as error by our filterwarnings) when the result isn't a
     # multiple of the zarr target's chunks. da.store writes blocks 1:1.
-    da.store(target_array, target, lock=False)  # type: ignore
+    # The shared lock serialises chunk read-modify-writes: for sharded
+    # targets the blocks (aligned to inner chunks) share shard objects and
+    # would otherwise race and silently drop data.
+    da.store(target_array, target, lock=DASK_STORE_LOCK)  # type: ignore
 
 
 def _on_disk_coarsen(
@@ -97,8 +101,8 @@ def _on_disk_coarsen(
         aggregation_function, source_array, coarsening_setup, trim_excess=True
     )
     out_target = out_target.rechunk(target.chunks)
-    # See _on_disk_dask_zoom for rationale.
-    da.store(out_target, target, lock=False)  # type: ignore
+    # See _on_disk_dask_zoom for rationale, including the lock.
+    da.store(out_target, target, lock=DASK_STORE_LOCK)  # type: ignore
 
 
 def on_disk_zoom(

src/ngio/images/_abstract_image.py

@@ -176,6 +176,15 @@ def chunks(self) -> tuple[int, ...]:
         """Return the chunks of the image."""
         return self.zarr_array.chunks
 
+    @property
+    def write_granularity(self) -> tuple[int, ...]:
+        """Return the atomic write unit of the on-disk array.
+
+        The shard shape when the array is sharded (writes are read-modify-writes
+        of whole shard objects), otherwise the chunk shape.
+        """
+        return self.zarr_array.shards or self.zarr_array.chunks
+
     @property
     def is_3d(self) -> bool:
         """Return True if the image is 3D."""

src/ngio/io_pipes/__init__.py

@@ -49,9 +49,11 @@
 )
 from ngio.io_pipes._match_shape import dask_match_shape, numpy_match_shape
 from ngio.io_pipes._ops_slices import SlicingInputType, SlicingOps, SlicingType
+from ngio.io_pipes._ops_slices_utils import ChunkRect
 from ngio.io_pipes._ops_transforms import TransformProtocol
 
 __all__ = [
+    "ChunkRect",
     "DaskGetter",
     "DaskGetterMasked",
     "DaskRoiGetter",

src/ngio/io_pipes/_ops_slices.py

@@ -10,6 +10,7 @@
 from pydantic import BaseModel, ConfigDict
 
 from ngio.common._dimensions import Dimensions
+from ngio.common._locks import DASK_STORE_LOCK
 from ngio.io_pipes._ops_slices_utils import compute_slice_chunks
 from ngio.ome_zarr_meta.ngio_specs import Axis
 from ngio.utils import NgioUserWarning, NgioValueError
@@ -255,7 +256,10 @@ def set_slice_as_dask(
         # da.store instead of da.to_zarr: see ngio.common._pyramid for the
         # dask>=2025.11 PerformanceWarning regression that to_zarr triggers
         # when the input chunks aren't a multiple of the target's chunks.
-        da.store(patch, zarr_array, regions=slice_tuple, lock=False)  # type: ignore
+        # The shared lock serialises chunk read-modify-writes: blocks
+        # misaligned with the target chunk (or shard) grid would otherwise
+        # race and silently drop data.
+        da.store(patch, zarr_array, regions=slice_tuple, lock=DASK_STORE_LOCK)  # type: ignore
         return
 
     # Complex case, we have exactly one tuple in the slicing tuple
@@ -264,7 +268,7 @@ def set_slice_as_dask(
         _sub_slice = (*slice_tuple[:ax], slice(idx, idx + 1), *slice_tuple[ax + 1 :])
         sub_patch = da.take(patch, indices=i, axis=ax)
         sub_patch = da.expand_dims(sub_patch, axis=ax)
-        da.store(sub_patch, zarr_array, regions=_sub_slice, lock=False)  # type: ignore
+        da.store(sub_patch, zarr_array, regions=_sub_slice, lock=DASK_STORE_LOCK)  # type: ignore
 
 
 ##############################################################

src/ngio/io_pipes/_ops_slices_utils.py

@@ -176,25 +176,112 @@ def compute_slice_chunks(
     return {tuple(idx) for idx in product(*per_axis_chunks)}
 
 
+ChunkRect: TypeAlias = tuple[tuple[int, int], ...]
+"""Per-axis inclusive `(first, last)` chunk-index ranges.
+
+Empty selections are represented as `None` wherever a `ChunkRect` is expected,
+never as a `first > last` range: an empty selection touches no chunks and must
+conflict with nothing.
+"""
+
+
+def _chunk_range_for_axis(
+    sel: SlicingType, size: int, csize: int
+) -> tuple[int, int] | None:
+    """The inclusive (first, last) chunk-index range touched on a single axis.
+
+    Exact for `slice` and `int` selections (their chunk footprint is contiguous,
+    see `_chunk_indices_for_axis`). For `list[int]` selections this is the
+    bounding box — a conservative over-approximation. Returns `None` for an
+    empty selection.
+    """
+    if isinstance(sel, slice):
+        start, stop = _normalize_slice(sel, size)
+        if start >= stop:  # empty
+            return None
+        return start // csize, (stop - 1) // csize
+
+    if isinstance(sel, int):
+        if sel < 0 or sel >= size:
+            raise IndexError(f"index {sel} out of bounds for axis of size {size}")
+        return sel // csize, sel // csize
+
+    if isinstance(sel, list):
+        if not sel:
+            return None
+        for v in sel:
+            if not isinstance(v, int):
+                raise TypeError("Only integers allowed inside tuple selections")
+            if v < 0 or v >= size:
+                raise IndexError(f"index {v} out of bounds for axis of size {size}")
+        return min(sel) // csize, max(sel) // csize
+
+    raise TypeError(f"Unsupported index type: {type(sel)!r}")
+
+
+def compute_chunk_rect(
+    shape: tuple[int, ...],
+    chunks: tuple[int, ...],
+    slicing_tuple: tuple[SlicingType, ...],
+) -> ChunkRect | None:
+    """Bounding rectangle, in chunk-index space, of the chunks a selection touches.
+
+    Exact for `slice`/`int` selections; for `list[int]` selections the per-axis
+    bounding box is used — a conservative over-approximation that may include
+    chunks the selection does not touch, never fewer.
+
+    Args:
+        shape: Overall array shape.
+        chunks: Chunk (or shard) shape — the conflict granularity.
+        slicing_tuple: Tuple of slices, ints, or lists of ints.
+
+    Returns:
+        Per-axis inclusive `(first, last)` chunk-index ranges, or `None` when
+        the selection is empty on any axis (touches no chunks).
+    """
+    if len(slicing_tuple) != len(shape):
+        raise NgioValueError(
+            f"key must have {len(shape)} items, got {len(slicing_tuple)}"
+        )
+
+    per_axis_ranges = []
+    for sel, size, csize in zip(slicing_tuple, shape, chunks, strict=True):
+        ax_range = _chunk_range_for_axis(sel, size, csize)
+        if ax_range is None:
+            return None
+        per_axis_ranges.append(ax_range)
+    return tuple(per_axis_ranges)
+
+
+def chunk_rects_intersect(rect_a: ChunkRect | None, rect_b: ChunkRect | None) -> bool:
+    """Whether two chunk rectangles share at least one chunk.
+
+    `None` encodes an empty selection and never intersects anything.
+    """
+    if rect_a is None or rect_b is None:
+        return False
+    if len(rect_a) != len(rect_b):
+        raise NgioValueError("Chunk rectangles must have the same rank")
+    return all(
+        a_first <= b_last and b_first <= a_last
+        for (a_first, a_last), (b_first, b_last) in zip(rect_a, rect_b, strict=True)
+    )
+
+
 def check_if_chunks_overlap(
     slices: Iterable[tuple[SlicingType, ...]],
     shape: tuple[int, ...],
     chunks: tuple[int, ...],
 ) -> bool:
-    """Check for overlaps in a list of slicing tuples using brute-force method.
+    """Check whether any two slicing tuples touch a common chunk.
+
+    This is O(n^2) in the number of slicing tuples, but each pair costs only a
+    per-axis range comparison — nothing is materialised. `list[int]` selections
+    are conservatively bounding-boxed: overlap may be reported for lists whose
+    exact chunk sets are disjoint, never the reverse.
 
-    This is O(n^2) and not efficient for large lists.
     Returns True if any overlaps are found.
     """
-    slices_chunks = (compute_slice_chunks(shape, chunks, si) for si in slices)
-    for it, (si, sj) in enumerate(_pairs_stream(slices_chunks)):
-        if si & sj:
-            return True
-        if it == 10_000:
-            warnings.warn(
-                "Performance Warning check_for_chunks_overlaps is O(n^2) and may be "
-                "slow for large numbers of regions.",
-                NgioUserWarning,
-                stacklevel=2,
-            )
-    return False
+    rects = (compute_chunk_rect(shape, chunks, si) for si in slices)
+    non_empty = (rect for rect in rects if rect is not None)
+    return any(chunk_rects_intersect(ri, rj) for ri, rj in _pairs_stream(non_empty))

src/ngio/iterators/__init__.py

@@ -2,7 +2,7 @@
 
 from ngio.iterators._feature import FeatureExtractorIterator
 from ngio.iterators._image_processing import ImageProcessingIterator
-from ngio.iterators._mappers import BasicMapper, MapperProtocol
+from ngio.iterators._mappers import BasicMapper, IterUnit, MapperProtocol
 from ngio.iterators._segmentation import (
     MaskedSegmentationIterator,
     SegmentationIterator,
@@ -12,6 +12,7 @@
     "BasicMapper",
     "FeatureExtractorIterator",
     "ImageProcessingIterator",
+    "IterUnit",
     "MapperProtocol",
     "MaskedSegmentationIterator",
     "SegmentationIterator",