Plan conflict-free write batches for mismatched chunk grids

.github/workflows/zarr-indexing.yml

@@ -43,9 +43,9 @@ jobs:
         uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4
       - name: Set up Python ${{ matrix.python-version }}
         run: uv python install ${{ matrix.python-version }}
-      # The suite imports nothing from `zarr`; it runs against the repo-root
-      # environment, with this package as an editable overlay, to share the
-      # parent project's pinned test toolchain. The recipes carry that
+      # The suite runs against the repo-root environment for optional Zarr
+      # codec integration and the parent's pinned toolchain, with this package
+      # as an editable overlay. The recipes carry that
       # invocation; this step only fixes the interpreter the matrix asked for.
       - name: Sync test dependency group
         run: uv sync --project ../.. --group test --python ${{ matrix.python-version }}

packages/zarr-indexing/benchmarks/README.md

@@ -26,3 +26,29 @@ These scripts measure planning rather than codec or storage throughput. Repeat
 measurements with alternating operation order before interpreting small timing
 differences. Preserve raw benchmark output as an experiment artifact rather than
 accumulating successive result tables in this README.
+
+## Write scheduling versus transfer planning
+
+`write_scheduling.py` compares structural outputs on aligned, misaligned,
+row-to-column, coarse write-unit, and hot-unit workloads:
+
+```sh
+hatch run test.py3.12-minimal:python packages/zarr-indexing/benchmarks/write_scheduling.py --planner native
+hatch run test.py3.12-minimal:python packages/zarr-indexing/benchmarks/write_scheduling.py --planner dask
+hatch run test.py3.12-minimal:python packages/zarr-indexing/benchmarks/write_scheduling.py --planner rechunker
+```
+
+Dask and Rechunker are optional comparison dependencies, not package runtime or
+CI dependencies. Use separate environments if their Zarr requirements conflict.
+Each invocation emits package version, geometry, timing, allocation, and either
+native schedule counts or the external planner's transfer stages. Record the
+checkout revision alongside the output. External planners receive a block memory
+budget of four times the larger source/target chunk's float64 byte size; the
+native scheduler makes no byte-memory guarantee and preserves source tasks.
+
+The Dask measurement calls its task planner without constructing or running the
+whole graph. Rechunker's result contains read/intermediate/write block shapes.
+Native results enumerate actual task assignments. Their constructor times are
+therefore not like-for-like performance scores. Use the comparison to identify
+when changing task boundaries is preferable to serializing existing tasks.
+See [the guide](../docs/guide/write-batches.md) for contracts and prior-art links.

packages/zarr-indexing/benchmarks/write_scheduling.py

@@ -0,0 +1,121 @@
+"""Structural planner comparison; separate planner runs may use separate environments.
+
+No data is copied. Native returns write batches; Dask and Rechunker return
+transfer stages with different execution/memory contracts. Times are not
+end-to-end rechunk throughput. --planner rechunker requires Rechunker separately.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import importlib.metadata
+import json
+import math
+import statistics
+import time
+import tracemalloc
+from typing import Any
+
+CASES = [
+    ("aligned", (120, 120), (12, 12), (12, 12)),
+    ("misaligned", (120, 120), (12, 12), (16, 16)),
+    ("row_to_column", (128, 128), (1, 128), (128, 1)),
+    ("shard_units", (120, 120), (12, 12), (60, 60)),
+    *[(f"hot_unit_{n}", (n,), (1,), (n,)) for n in (100, 1000, 10000)],
+]
+
+
+def split(shape: tuple[int, ...], chunks: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
+    return tuple(
+        tuple(min(c, n - i) for i in range(0, n, c)) for n, c in zip(shape, chunks, strict=True)
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--planner", choices=("native", "dask", "rechunker"), default="native")
+    parser.add_argument("--repeats", type=int, default=3)
+    args = parser.parse_args()
+    if args.repeats < 1:
+        parser.error("repeats must be positive")
+    results: dict[str, Any] = {}
+    for name, shape, source, target in CASES:
+        max_mem = 4 * 8 * max(math.prod(source), math.prod(target))
+
+        def operation(
+            shape: tuple[int, ...] = shape,
+            source: tuple[int, ...] = source,
+            target: tuple[int, ...] = target,
+            max_mem: int = max_mem,
+        ) -> Any:
+            if args.planner == "native":
+                from zarr_indexing import IndexDomain, plan_rechunk
+                from zarr_indexing.grid import dimension_grids_from_chunks
+
+                return plan_rechunk(
+                    IndexDomain.from_shape(shape),
+                    dimension_grids_from_chunks(source, shape),
+                    dimension_grids_from_chunks(target, shape),
+                )
+            if args.planner == "dask":
+                module = importlib.import_module("dask.array.rechunk")
+                return module.plan_rechunk(
+                    split(shape, source), split(shape, target), itemsize=8, block_size_limit=max_mem
+                )
+            from rechunker.algorithm import rechunking_plan
+
+            return rechunking_plan(shape, source, target, itemsize=8, max_mem=max_mem)
+
+        operation()  # imports and warmup excluded
+        samples = []
+        for _ in range(args.repeats):
+            start = time.perf_counter()
+            plan = operation()
+            samples.append((time.perf_counter() - start) * 1000)
+        tracemalloc.start()
+        allocated = operation()
+        _, peak = tracemalloc.get_traced_memory()
+        tracemalloc.stop()
+        del allocated
+        row: dict[str, Any] = {
+            "shape": shape,
+            "source": source,
+            "destination_units": target,
+            "median_ms": statistics.median(samples),
+            "peak_mib": peak / 2**20,
+        }
+        if args.planner == "native":
+            row.update(
+                tasks=len(plan.pieces),
+                batches=len(plan.schedule.batches),
+                max_parallel_tasks=max(map(len, plan.schedule.batches), default=0),
+                write_units=plan.schedule.n_write_units,
+                task_unit_memberships=plan.schedule.n_memberships,
+            )
+        elif args.planner == "dask":
+            row.update(max_mem=max_mem, transfer_stages=plan)
+        else:
+            read_chunks, intermediate_chunks, write_chunks = plan
+            row.update(
+                max_mem=max_mem,
+                read_chunks=read_chunks,
+                intermediate_chunks=intermediate_chunks,
+                write_chunks=write_chunks,
+            )
+        results[name] = row
+    package = "zarr-indexing" if args.planner == "native" else args.planner
+    print(
+        json.dumps(
+            {
+                "planner": args.planner,
+                "version": importlib.metadata.version(package),
+                "results": results,
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

packages/zarr-indexing/changes/324.feature.md

@@ -0,0 +1,6 @@
+Add `plan_write_batches` to group transformed writes into batches with disjoint
+destination write units, with explicit preserve-order and reorder policies.
+`plan_rechunk` keeps each source chunk as a task when copying between different
+grids, including irregular layouts and clipped domains. Plans describe work
+without executing I/O; callers supply the storage write-unit grid and enforce
+completion barriers between batches.

packages/zarr-indexing/docs/api/index.md

@@ -34,6 +34,10 @@ and the wire format built on top of it.
 
 **Chunk resolution**
 
+- [`zarr_indexing.scheduling`](scheduling.md) — `plan_write_batches` and
+  `plan_rechunk`, metadata-only schedules for conflicting destination write units;
+  start with [Conflict-free write batches](../guide/write-batches.md).
+
 - [`zarr_indexing.chunk_resolution`](chunk_resolution.md) —
   `plan_chunks`, which lazily projects a request through a caller-selected grid,
   the reusable `ChunkPlan` and paired-transform `ChunkProjection` values, and

packages/zarr-indexing/docs/api/scheduling.md

@@ -0,0 +1,5 @@
+---
+title: scheduling
+---
+
+::: zarr_indexing.scheduling

packages/zarr-indexing/docs/guide/integrations.md

@@ -1,5 +1,8 @@
 # Integration boundaries
 
+For incoming chunks that do not align with the destination's independently
+writable units, see [Conflict-free write batches](write-batches.md).
+
 For the complete path from indexing syntax to chunk coordinates, local selectors,
 and result positions, start with [From a selection to chunk operations](selection-flow.md).
 
@@ -37,8 +40,8 @@ prefetcher — does not need a `ChunkProjection` object per chunk. The plan's
 [factored form](index.md#a-plan-is-a-product-of-per-axis-tables) is a few
 NumPy arrays per axis, and everything a chunk copy needs is a row of each:
 the chunk index, the chunk-local start and extent, and where the cells land
-in the request. `chunk_coords()` alone answers "which chunks?" for a prefetch, without
-materializing anything.
+in the request. `chunk_coords()` answers "which chunks?" by allocating their coordinate
+array without constructing projections.
 
 The example assembles a strided box from its `StridedSet` tables, one slice
 per chunk. Three things a real consumer also has to get right are checked at

packages/zarr-indexing/docs/guide/write-batches.md

@@ -0,0 +1,147 @@
+# Conflict-free write batches
+
+Data often arrives in chunks that do not align with the destination chunks.
+Disjoint array slices can still race: both writes may read, modify, and replace
+the same destination chunk. `plan_rechunk` keeps each incoming source chunk as
+one task and groups tasks into batches with disjoint destination write units.
+It plans the work; your code reads buffers, submits tasks, and writes data.
+
+## Source chunks stay intact
+
+A length-12 source arriving in chunks of 3 and targeting chunks of 4 produces:
+
+| Source task | Global values | Destination units | Batch |
+| --- | --- | --- | --- |
+| 0 | `[0:3]` | `{0}` | 0 |
+| 1 | `[3:6]` | `{0, 1}` | 1 |
+| 2 | `[6:9]` | `{1, 2}` | 0 |
+| 3 | `[9:12]` | `{2}` | 1 |
+
+Each source chunk is one task even when it touches several destination units.
+Tasks 0 and 2 run concurrently, then tasks 1 and 3. Reordering is valid because
+their logical destination elements are disjoint. Chunk boundaries still require
+read-modify-write, so all four tasks cannot safely run at once.
+
+```text
+source chunks + global domain
+       -> one piece per intersecting source chunk
+       -> destination transform per piece
+       -> touched destination write-unit coordinates
+       -> batches with no shared write unit
+       -> caller executes each batch, then waits for completion
+```
+
+A `RechunkPiece` has a source chunk coordinate, slices relative to that decoded
+source chunk, and a destination transform. The transform is an identity over
+the piece's global domain; its bounds give the destination slices. A subdomain
+clips source selections without changing chunk identity. Irregular grids,
+nonzero origins, empty domains, and scalar arrays are supported.
+
+The example uses already-arrived NumPy chunks and a NumPy destination to expose
+the task boundary. Replace `target[selection] = values` with the destination's
+write operation. Real Zarr chunked and sharded writes are covered by the package's
+codec integration tests.
+
+```python
+--8<-- "snippets/write_batches.py:copy"
+```
+
+## What makes a batch safe
+
+Supply the grid of **independently writable storage units**. For sharded Zarr,
+that is generally the shard grid rather than the inner codec-chunk grid. When
+metadata or an enclosing object couples otherwise separate chunks, use the
+coarser transaction unit, or provide external coordination.
+
+Complete all destination reads and writes in batch *n* before starting any
+destination read-modify-write for batch *n + 1*. Merely submitting futures is
+not a barrier. Prefetching destination snapshots across batches can lose updates
+even with a correct schedule. Independent source data may be prefetched.
+
+The source and destination must not alias during a copy. The schedule knows
+write conflicts, not read-after-write dependencies between source and target.
+It does not coordinate unrelated writers, retries still running after a failed
+batch, metadata changes, or mutations to the supplied grids. Each task must also
+perform its own writes safely; the schedule coordinates between tasks.
+
+## Scheduling existing transformed tasks
+
+`plan_write_batches(writes, destination_grid)` accepts a finite iterable of
+request-to-destination `IndexTransform`s. Task IDs are their zero-based input
+positions. Keep your buffers or futures separately and use these IDs to find
+them. Every task appears once, including an empty write; an empty iterable
+produces no batches. Duplicate coordinates within one task occupy a unit once
+for scheduling, but assignment semantics within that task remain your concern.
+
+The default `order="preserve"` retains input order between tasks that touch any
+of the same destination units. Unrelated tasks can move earlier; this is not a
+guarantee about arbitrary side effects. `order="reorder"` explicitly permits
+reordering and uses deterministic first-fit coloring. For overlapping logical
+writes, that can change the final values. `plan_rechunk` uses reorder because
+its identity-copy pieces are logically disjoint.
+
+Footprints come from the existing chunk planner, preserving its supported
+transform semantics. Affine diagonals and mixed affine/index-array maps sharing
+an input axis are unsupported by the factored planner and raise. Nonempty out-of-bounds transforms
+raise before a schedule is returned. See the [scheduling API](../api/scheduling.md).
+
+## Costs and limits
+
+Preparation is eager. The returned schedule contains immutable tuples of task
+IDs; a rechunk plan additionally retains piece metadata. It holds no array data.
+It does not enforce a byte-memory budget, split oversized source chunks, choose
+temporary storage, or guarantee the fewest batches.
+
+Preserve mode tracks the last batch using each destination unit. Reorder mode
+tracks occupied colors and the first free color per unit; this avoids repeatedly
+scanning a hot unit's full history. It still uses a heuristic with no linear
+worst-case runtime claim. No pairwise conflict graph is built. Planning memory
+includes the current footprint, chunk-planner intermediates, task IDs, and
+per-unit state. Reorder state scales with task/unit memberships; large footprints
+and many source tasks can still consume substantial memory.
+
+`n_write_units` counts distinct destination units and `n_memberships` sums the
+units touched per task. They are structural diagnostics, not exact byte-I/O or
+peak-memory predictions. To limit execution memory, independently cap concurrent
+tasks inside a batch. Splitting a batch into smaller sequential sub-batches
+preserves safety; combining batches does not necessarily do so.
+
+A row-to-column transfer exposes the limitation: every row task touches every
+column write unit, so rows must run serially. No recoloring can create parallelism
+while preserving those tasks. Destination-owned assembly or an intermediate
+layout changes the task boundaries and is often the better strategy.
+
+## Comparison with existing planners
+
+These systems solve related but different problems:
+
+| Planner | Work it chooses | Main constraint | Execution consequence |
+| --- | --- | --- | --- |
+| This utility | Batches of existing source tasks | Disjoint destination write units per batch | Repeated destination RMW and serialization are possible |
+| Rechunker | Consolidated read/write regions, optionally via an intermediate layout | Configured worker memory and reduced transfer overhead | Can change task boundaries and use temporary storage |
+| Dask task rechunking | Split/merge transfer stages between chunk layouts | Graph growth and block-size limits | Constructs a transfer graph rather than storage-write batches |
+
+[Rechunker's algorithm](https://rechunker.readthedocs.io/en/stable/algorithm.html)
+selects consolidated read and write chunks and uses an intermediate array when
+needed. Its [planner implementation](https://github.com/pangeo-data/rechunker/blob/v0.5.4/rechunker/algorithm.py)
+also includes multistage planning. It is a better fit when transfer layout and
+temporary storage can be chosen under a memory budget.
+
+[Dask's task planner](https://docs.dask.org/en/stable/_modules/dask/array/rechunk.html)
+chooses intermediate layouts using graph-size and block-size controls. Dask also
+has a peer-to-peer rechunk path; the comparison script exercises only its task
+planner, not that distributed execution path.
+
+[Xarray's `to_zarr`](https://docs.xarray.dev/en/stable/generated/xarray.Dataset.to_zarr.html)
+checks chunk alignment for parallel safety and provides alignment options.
+This scheduler offers a different tradeoff for fixed incoming tasks: keep their
+boundaries and serialize conflicts. It does not disable or replace Xarray's
+checks automatically.
+
+`benchmarks/write_scheduling.py` runs each planner separately, recording native
+batch/task/membership counts, Dask's transfer layouts, or Rechunker's named
+read/intermediate/write block shapes. Those three shapes do not imply three
+execution stages. External
+planners can run in their own environments, avoiding dependency conflicts.
+Planning times describe different returned products and are not interchangeable
+performance scores or end-to-end copy measurements.

packages/zarr-indexing/docs/snippets/write_batches.py

@@ -0,0 +1,37 @@
+"""Schedule already chunked data without changing the incoming chunk tasks."""
+
+# --8<-- [start:copy]
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+
+from zarr_indexing import IndexDomain, plan_rechunk
+from zarr_indexing.grid import dimension_grids_from_chunks
+
+source_chunks = {i: np.arange(3 * i, 3 * i + 3) for i in range(4)}
+target = np.full(12, -1)
+plan = plan_rechunk(
+    IndexDomain.from_shape((12,)),
+    dimension_grids_from_chunks((3,), (12,)),
+    dimension_grids_from_chunks((4,), (12,)),
+)
+assert plan.schedule.batches == ((0, 2), (1, 3))
+
+
+def copy_piece(task: int) -> None:
+    piece = plan.pieces[task]
+    values = source_chunks[piece.source_chunk[0]][piece.source_selection]
+    domain = piece.destination.domain
+    selection = tuple(
+        slice(lo, hi) for lo, hi in zip(domain.origin, domain.exclusive_max, strict=True)
+    )
+    target[selection] = values
+
+
+with ThreadPoolExecutor(max_workers=4) as pool:
+    for batch in plan.schedule.batches:
+        # Consume results and propagate errors before the next batch starts.
+        list(pool.map(copy_piece, batch))
+
+np.testing.assert_array_equal(target, np.arange(12))
+# --8<-- [end:copy]