Coastal & Marine Spatial Analysis Pipelines

Rechunking a Zarr Store Without a Full Rewrite

A store chunked for map reads is being hammered by time-series queries, and the fix is to change the chunk shape. That operation is a full shuffle: every output chunk draws from many input chunks, so the naive implementation — read the array, rechunk, write — needs the whole thing in memory or spills catastrophically. This page is the staged alternative and the cost model behind it, within the Zarr versus NetCDF topic.

Root cause: rechunking is an all-to-all data movement

Consider a store of shape (8760, 4096, 4096) chunked (1, 512, 512) — one timestep per chunk, good for maps. Rechunking to (8760, 32, 32) for time series means every output chunk needs one 32×32 patch from all 8,760 timesteps, and each of those patches lives in a different input chunk. The output chunk therefore touches 8,760 input chunks; the input chunk contributes to 256 output chunks.

import numpy as np

def shuffle_factor(shape, src_chunks, dst_chunks) -> float:
    """How many input chunks a typical output chunk touches."""
    if len(shape) != len(src_chunks) != len(dst_chunks):
        raise ValueError("shape and both chunk tuples must have the same rank")
    touched = 1.0
    for s, d in zip(src_chunks, dst_chunks):
        touched *= np.ceil(d / s) if d > s else 1.0
    return float(touched)

print(shuffle_factor((8760, 4096, 4096), (1, 512, 512), (8760, 32, 32)))   # 8760.0
print(shuffle_factor((8760, 4096, 4096), (1, 512, 512), (24, 512, 512)))   # 24.0

That number is the whole cost model. Dask’s default rechunk builds a task for each input-output pair, so a shuffle factor of 8,760 over 6,000 output chunks is fifty million tasks — a graph the scheduler cannot hold, let alone execute, as profiling Dask graph overhead sets out.

One All-to-All Shuffle, or Two Bounded Ones Two graphs of the same rechunk. On the left, a direct rechunk draws a dense mesh of lines from six input chunks to six output chunks, every input connected to every output, labelled as an all-to-all shuffle whose peak memory is the full array because no output chunk can be finalised until all its inputs have been read. On the right, the same rechunk goes through an intermediate chunk shape: six inputs connect to four intermediate chunks in a sparse pattern, and those four connect to six outputs, also sparsely. A note records that the staged version moves more bytes in total but never holds more than one intermediate chunk group in memory, which is the property that makes it finish. direct — all to all peak memory = the whole array the graph does not fit either staged through an intermediate more bytes moved in total peak memory is a number you choose The staged version is slower on paper and is the one that finishes.

Step-by-step with production code

1. Confirm the current chunking really is the problem

Rechunking a terabyte store costs real money, and the read pattern that motivated it is often served better by a small change elsewhere. Measure before committing:

import logging
import time

import xarray as xr

logger = logging.getLogger(__name__)


def time_read_patterns(store: str) -> dict[str, float]:
    """Time the two access patterns that chunk shape trades between."""
    ds = xr.open_zarr(store, consolidated=True)
    var = next(iter(ds.data_vars))

    t0 = time.perf_counter()
    ds[var].isel(time=0).load()
    map_s = time.perf_counter() - t0

    t0 = time.perf_counter()
    ds[var].isel(y=len(ds.y) // 2, x=len(ds.x) // 2).load()
    series_s = time.perf_counter() - t0

    logger.info("one map %.2f s · one time series %.2f s · ratio %.0f×",
                map_s, series_s, series_s / max(map_s, 1e-9))
    return {"map_s": map_s, "series_s": series_s}

A ratio under about twenty is usually not worth a rechunk; caching, a coarser overview or simply issuing the reads concurrently will close it. Above a hundred, the shape is genuinely wrong for the workload.

2. Stage the shuffle with rechunker

rechunker plans the intermediate shape from a memory budget you supply, writes to a temporary store, and then writes the target. Memory is bounded by the budget rather than by the array:

def rechunk_store(src: str, dst: str, tmp: str, target_chunks: dict,
                  max_mem: str = "2GB") -> None:
    """Staged rechunk with bounded memory. `tmp` must be deleted afterwards."""
    import zarr
    from rechunker import rechunk

    source = zarr.open_consolidated(src, mode="r")
    plan = rechunk(source, target_chunks=target_chunks, max_mem=max_mem,
                   target_store=dst, temp_store=tmp)
    logger.info("rechunk plan: %s", plan)
    plan.execute()
    zarr.consolidate_metadata(dst)
    logger.info("wrote %s and consolidated its metadata", dst)

Two things bite here. The temporary store is roughly the size of the array and is not cleaned up automatically — on object storage that is a real bill accruing until someone deletes it, and a lifecycle rule on the temp prefix is the reliable fix. And the target store needs its own consolidate_metadata call; the source’s .zmetadata describes the source’s chunk shape, and copying it produces a store whose metadata lies about its own layout, which is the failure the consolidated-metadata guide covers.

3. Consider two views instead of one rechunk

A rechunk exists to serve a second read pattern. Publishing a second store serves it too, and the arithmetic is often favourable: a rechunk costs one full read plus two full writes and leaves one store, while a second view costs one full read plus one full write and leaves two stores that are each optimal.

The trade is storage against compromise. Where the store is small relative to the budget, two views is strictly better — each reader gets an exact fit and neither pays for the other. Where storage dominates, or where the two views would drift out of sync as data is appended, one store and a compromise chunking is the right answer.

Rechunk, or Publish Twice A two-column comparison. The rechunk column shows one full read of the source, a full write to the temporary store, a full read back and a full write to the target, ending with one store in the new shape; the map reader now pays a penalty because the new shape favours time series. The second-view column shows one full read and one full write, ending with two stores, one chunked for maps and one for time series, each serving its reader exactly. A footer notes that the second view costs more storage and less compute, that it is strictly better when storage is cheap relative to the read volume, and that it becomes harder to maintain once the store is appended to, because both views must be updated together. rechunk in place read source in full write temp store in full read temp in full write target in full result: 1 store, new shape the map reader now pays instead temp store must be deleted publish a second view read source in full write second store in full result: 2 stores maps served exactly · series served exactly 2× storage both must be updated on append Half the I/O and no compromise, for twice the storage. On object storage that trade usually favours two views — until the store starts being appended to, when keeping them in step becomes the dominant cost.

Verification and acceptance test 

import numpy as np
import pytest


def test_rechunk_preserved_the_values(src: str, dst: str):
    a = xr.open_zarr(src, consolidated=True)
    b = xr.open_zarr(dst, consolidated=True)
    assert a.dims == b.dims, f"dims changed: {a.dims} vs {b.dims}"

    rng = np.random.default_rng(0)
    for _ in range(20):
        idx = {d: int(rng.integers(0, a.sizes[d])) for d in a.dims}
        for v in a.data_vars:
            x, y = float(a[v].isel(idx)), float(b[v].isel(idx))
            assert np.isclose(x, y, equal_nan=True), f"{v} at {idx}: {x} vs {y}"


def test_the_new_store_has_its_own_consolidated_metadata(dst: str):
    import zarr
    z = zarr.open_consolidated(dst, mode="r")
    for name, arr in z.arrays():
        assert arr.chunks == TARGET_CHUNKS[name], (
            f"{name} reports chunks {arr.chunks}; the metadata is the old snapshot")

Sampling random indices rather than comparing shapes is the point of the first test: a shuffle that transposed two axes produces a store of the correct shape holding the wrong values everywhere, and only a value comparison catches it.

A Rechunk Moves the Cost; It Does Not Remove It Two paired bars comparing read times before and after a rechunk from map-optimised to series-optimised chunking. Before, a single map read takes 0.4 seconds and a single point time series takes 47 seconds. After, the map read takes 6 seconds and the time series takes 0.3 seconds. A weighted-total panel shows the combined daily cost under two workload mixes: one dominated by map reads, where the rechunk makes the system slower overall, and one dominated by time-series reads, where it makes it far faster. A caption states that the rechunk is justified only by the measured mix and never by the individual timings. one read, before and after map read before 0.4 s after 6 s — 15× slower time series before 47 s after 0.3 s — 150× faster daily total, under two workload mixes 90% maps rechunk is a loss 90% series a large win The individual timings justify nothing. Only the measured mix decides, which is why step one is a measurement.

Edge cases and gotchas

  • The temp store outlives the job. rechunker does not delete it, and a failed run leaves it behind. Put it under a prefix with a short lifecycle rule, as in setting S3 lifecycle rules, rather than relying on a cleanup step that only runs on success.
  • Compression interacts with chunk size. A larger chunk usually compresses better, so the on-disk size after a rechunk is not simply the same bytes rearranged. Budget for the target store being 10–30% smaller or larger, and measure rather than assume.
  • Appending after a rechunk needs the same shape. A store rechunked to (8760, 32, 32) cannot have a single new timestep appended cheaply, because every chunk spans the whole time axis. Time-series chunking and append-friendliness are in direct opposition, and a store that must accept daily appends should not be chunked along its full time extent.

Up: Zarr vs NetCDF for Cloud-Native Coastal Pipelines