Rechunker

Rechunker uses the Push / Pull Consolidated algorithm to rechunk datasets while satisfying three primary constraints: respecting memory limits, minimizing the number of required tasks (always less than N source chunks + M target chunks), and ensuring the process is embarrassingly parallel (avoiding write locks and inter-worker communication).

The algorithm follows these steps:

1. Input Configuration: The user provides a source array (shape, chunk structure, and data type), the desired target_chunks structure, and a max_mem threshold.
2. Determine Write Batches: The algorithm calculates the largest batch of data a single worker can write without exceeding max_mem. These are called write_chunks.
3. Determine Read Batches: The algorithm calculates the largest batch of data a single worker can read given max_mem, attempting to align these with write_chunks where possible. These are called read_chunks.
4. Direct Copy Path: If write_chunks == read_chunks, the algorithm performs a direct copy from the source to the target, avoiding intermediate datasets.
5. Intermediate Path: If they do not match, the algorithm defines intermediate chunks as the minimum of write_chunks and read_chunks along each axis. Data is first copied from the source to an intermediate array, and then from the intermediate array to the target.

Rechunker uses a staged approach to define execution workflows.

Stage

A Stage is a frozen dataclass representing a single step in a workflow. It contains:

• function: A StageFunction, which can be a function taking no arguments (except config) or a function taking a single config argument.
• name: A string identifier for the stage.
• mappable: An optional Iterable used for mapping operations.

Pipeline

A Pipeline is a frozen dataclass that groups a sequence of Stage objects together with an optional config object.

ParallelPipelines

A ParallelPipelines type is a Tuple of one or more Pipeline objects intended to be executed in parallel.

from rechunker.types import Stage, Pipeline

def my_stage_func(config=None):
    pass

stage = Stage(function=my_stage_func, name="my_stage")
pipeline = Pipeline(stages=[stage], config={"param": 1})
parallel_pipelines = (pipeline,)

To rechunk a single array, use the rechunk function. You must provide the source array, the desired target_chunks (as a tuple or a dictionary mapping dimension names to sizes), the max_mem allowed for intermediate steps, a target_store for the final data, and a temp_store for intermediate data.

Note that rechunk returns a Rechunked object, which represents a plan. The actual rechunking process is only performed when you call the .execute() method on this object. Rechunking can be performed using the default Dask multi-threaded scheduler or by providing a Dask Client (e.g., from a LocalCluster or GatewayCluster) to use a distributed cluster.

from rechunker import rechunk

target_chunks = {"time": 2920, "lat": 25, "lon": 1}
max_mem = "1MB"
target_store = "air_rechunked.zarr"
temp_store = "air_rechunked-tmp.zarr"

# Create the plan
array_plan = rechunk(
    source_array, 
    target_chunks, 
    max_mem, 
    target_store, 
    temp_store=temp_store
)

# Execute the plan
result = array_plan.execute()

Rechunker can be used to rechunk data stored in cloud buckets (like Google Cloud Storage) by using a filesystem mapper (e.g., gcsfs).

1. Use gcsfs.GCSFileSystem to create a filesystem object.
2. Use gcs.get_mapper(url) to create a mapper for the source and target locations.
3. Pass these mappers to the rechunk function as the target_store and temp_store.
4. For large-scale cloud rechunking, it is recommended to use a distributed Dask cluster (e.g., via Dask Gateway).

import gcsfs
from rechunker import rechunk
from dask.distributed import Client

# Setup cluster and filesystem
client = Client() # or a GatewayCluster
gcs = gcsfs.GCSFileSystem(requester_pays=True)

# Define locations using mappers
source_store = gcs.get_mapper("gs://pangeo-cmems-duacs")
target_store = gcs.get_mapper("gs://your-bucket/target.zarr")
temp_store = gcs.get_mapper("gs://your-bucket/temp.zarr")

# Open source array from group
import zarr
group = zarr.open_consolidated(source_store, mode="r")
source_array = group["sla"]

# Plan and execute
plan = rechunk(source_array, (8901, 72, 72), "1GB", target_store, temp_store=temp_store)
plan.execute()

Rechunker supports rechunking entire groups (to preserve metadata/coordinates) or Xarray Datasets.

When rechunking a group, target_chunks must be a dictionary where keys are variable names. For each variable, you can specify:

• A dictionary of dimension-to-chunksize mappings (e.g., {"time": 2920, "lat": 25, "lon": 1}).
• None to indicate that the specific array should not be rechunked.

If you provide a simple dictionary mapping dimension names to chunk sizes (e.g., {"time": 2920, "lat": 25, "lon": 1}) directly to the rechunk function when passing an Xarray Dataset, rechunker will apply those consistent chunks to all variables in the dataset.

# Rechunking a group with specific settings per variable
target_chunks = {
    "air": {"time": 2920, "lat": 25, "lon": 1},
    "time": None,  # don't rechunk this array
    "lon": None,
    "lat": None,
}

array_plan = rechunk(
    source_group, 
    target_chunks, 
    max_mem, 
    target_store, 
    temp_store=temp_store
)
array_plan.execute()