Best way of copying and rechunking a remote zarr store

[guidocioni]

I want to create a local copy of some of the variables contained in one of the many ERA5 zarr replicas available (here I chose to use the google cloud storage one for convenience). Why? Because we need quick access for analysis and I want to optimize the chunking in the process for our queries.

The idea is to initialize the skeleton and write the metadata to disk without computing

# Chunk specification
CHUNKS = {
    "time": 1,
    "level": 1,
    "latitude": 721,
    "longitude": 1440
}

for var_name in VARIABLES:
    da_var = ds_filtered[var_name]
    var_chunks = tuple([CHUNKS.get(d, da_var.sizes[d]) for d in da_var.dims])
    
    data_vars[var_name] = (
        da_var.dims,
        da.empty(da_var.shape, dtype=da_var.dtype, chunks=var_chunks)
    )

ds_empty = xr.Dataset(data_vars=data_vars, coords=coords, attrs=ds_filtered.attrs)

ds_empty.to_zarr(TARGET_ZARR, compute=False, consolidated=True, mode='w', zarr_format=2)

And then start writing with multiple workers in parallel to different regions

def process_time_chunk_like_api(time_index):
    global worker_ds
    
    ds_hour = worker_ds.isel(time=slice(time_index, time_index + 1))
    
    # Eagerly load the hour strictly into RAM.
    with dask.config.set(scheduler='single-threaded'):
        ds_hour = ds_hour.load()
            
    ds_hour = ds_hour.chunk(CHUNKS)
    
    ds_hour.to_zarr(TARGET_ZARR, region='auto', zarr_format=2)

This works but it’s not really scalable as it uses an insane amount of memory in the beginning and surely isn’t optimized. Even though the original gcs chunks contain all vertical levels in one, expanding them into memory shouldn’t cause the consumption I"m seeing.

I felt like this is something that rechunker was supposed to do, but I believe that library is outdated. Cubed does not seem to support multiple workers writing to the same zarr.

Surely I’m doing something wrong as this seems to be a recurring issue. What is the most efficient way of copying a zarr from one remote to another while rechunking and using the least amount of resources possible?

[aaronspring]

Hi Guido,

The init step doesn’t have to be expensive — make the template lazy and single-chunked, and set the target chunking through encoding rather than through the template’s own chunks:

import dask.array as da
import numpy as np
import xarray as xr

# one lazy chunk per variable: nothing is allocated, and the graph stays
# at one task per variable no matter how big the dataset is
data_vars = {
    name: xr.DataArray(
        da.full(shape, np.float32("nan"), chunks=-1, dtype="float32"),
        dims=dims,
        attrs=src[name].attrs,
    )
    for name, ... in ...
}
ds_empty = xr.Dataset(data_vars, coords=coords, attrs=attrs)

# the on-disk chunking is declared here, not by the template's dask chunks
# also set compression here
encoding = {v: {"chunks": target_chunks} for v in ds_empty.data_vars}

ds_empty.to_zarr(TARGET, compute=False, mode="w", encoding=encoding)

Keeps the deferred write graph tiny. If the template is numpy-backed, or chunked at the target chunk size, you either allocate the whole array or build one task per output chunk — which is where the memory blowup at init comes from. Good older thread on exactly this:

`.to_zarr(..., compute=False)` optimizations

The likely problem here is that the Dask graph for that array is absolutely huge, which bogs everything down. Instead, use this pattern: Just use one single Dask chunk for the Xarray dataset Specify chunk size via encoding Here’s an example of that from our Serverless Data Cube Demo: An even better solution would be to stop overloading to_zarr and actually implement schema creation in Xarray. That has been discussed extensively here:

Decoupling template chunks from store chunks via encoding means you can change the target layout without touching how the template is built.

Side questions:

• why using zarr_format=2?
• have you considered icechunk-era5 which comes with dual chunking to avoid having to download and rechunk yourself?

[guidocioni]

Hey @aaronspring

first some answers to your questions

• why zarr_format=2 → The original dataset is also zarr2 and when writing I had some issues due to different compression. I was too lazy to properly fix it so just decided to create an equivalent zarr2 for the moment.
• icecunk-era5 → Considered, yes, but I needed all vertical levels and variables. AFAIK the Google replica is the only fully complete ERA5 replica.

After some deeper exploration I realized the memory explosion was due to the usage of chunks={} when opening the zarr store. Due to the huge dimensionality of the gcs ERA5 zarr (200+ variables, “unlimited” time axis, 2D and 3D variables mixed together), xarray was spending most of the time at the beginning just reading the metadata and preparing everything.

Using chunks=None, which basically turns off Dask and reads everything sequentially, I avoided the initial memory blowup. This of course does not use the distributed engine, which means you also don’t see any progress when running a computation, but for a simple copy job it may not matter.

Some interesting benchmarks I ran while trying to just load a single time step of temperature

• chunks = {} → 31 seconds spent just opening the zarr, 20GB RAM usage, 34s computation time
• chunks = None → 2.9 seconds spent opening the zarr, 1minute 37s computation time, negligible RAM usage

This is the current version of my script (simplified) which works pretty well (I can run it with 30 concurrent processes and it only uses 50-60GB of RAM after the initial peak).

import xarray as xr
import dask.array as da
from concurrent.futures import ProcessPoolExecutor
from tqdm import tqdm

START_DATE = "2020-01-01"
END_DATE = "2026-07-01"

VARIABLES = [
    # 3D
    "temperature",
    # 2D
    "2m_temperature",
]

SOURCE_ZARR = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
TARGET_ZARR = "s3://zarr-export/era5_gcs_replica_test.zarr"

CHUNKS = {"time": 1, "level": 1, "latitude": 721, "longitude": 1440}
WORKERS = 30

worker_ds = None

def init_worker(source_zarr, start_date, end_date, variables):
    global worker_ds
    ds_src = xr.open_zarr(source_zarr, chunks=None, storage_options=dict(token="anon"))
    worker_ds = ds_src[variables].sel(time=slice(start_date, end_date))

def process_time_chunk(time_index):
    global worker_ds
    ds_hour = worker_ds.isel(time=slice(time_index, time_index + 1))
    ds_hour.load()

    for var in ds_hour.variables:
        ds_hour[var].encoding.pop("chunks", None)
        ds_hour[var].encoding.pop("preferred_chunks", None)

    ds_hour = ds_hour.chunk(CHUNKS)
    ds_hour.to_zarr(TARGET_ZARR, region="auto", zarr_format=2)

    return True

if __name__ == "__main__":
    ds_source = xr.open_zarr(
        SOURCE_ZARR, chunks=None, storage_options=dict(token="anon")
    )
    ds_filtered = ds_source[VARIABLES].sel(time=slice(START_DATE, END_DATE))

    coords = {dim: ds_filtered[dim] for dim in ds_filtered.dims}
    data_vars = {}

    for var_name in VARIABLES:
        da_var = ds_filtered[var_name]
        var_chunks = tuple([CHUNKS.get(d, da_var.sizes[d]) for d in da_var.dims])
        data_vars[var_name] = (
            da_var.dims,
            da.empty(da_var.shape, dtype=da_var.dtype, chunks=var_chunks),
        )

    ds_empty = xr.Dataset(data_vars=data_vars, coords=coords, attrs=ds_filtered.attrs)

    ds_empty.to_zarr(TARGET_ZARR, compute=False, mode="w", zarr_format=2)

    n_times = len(ds_filtered.time)
    init_args = (SOURCE_ZARR, START_DATE, END_DATE, VARIABLES)

    with ProcessPoolExecutor(
        max_workers=WORKERS, initializer=init_worker, initargs=init_args
    ) as executor:
        futures = [executor.submit(process_time_chunk, i) for i in range(n_times)]
        for future in tqdm(futures, total=n_times):
            future.result()

I believe this is still not ideal, as this seems to fail if I create a long time axis with empty values (like in the gcs replica), but I couldn’t really find a better way for the moment…