Tool Integration: Rechunker + Zarr + NetCDF — dask_setup Recipe
This notebook shows how dask_setup integrates with the broader scientific Python stack
for memory-efficient data conversion and storage:
Rechunker — safe NetCDF → Zarr conversion with a strict memory budget
Zarr — write optimally-chunked stores with metadata consolidation
Temporary directory management — using
dask_tmpfromsetup_dask_clientas the rechunker intermediate store (stays on fast local storage like$PBS_JOBFS)Validation — verify data integrity after conversion
Incremental Zarr writes — appending time windows to avoid OOM on large datasets
Setup
import shutil
import tempfile
import time
from pathlib import Path
import numpy as np
import xarray as xr
from dask_setup import rechunk_dataset, setup_dask_client
def blosc_encoding(cname="lz4", clevel=5):
"""Per-variable Zarr compressor encoding for whichever zarr major version is installed."""
if int(zarr.__version__.split(".")[0]) >= 3:
from zarr.codecs import BloscCodec
return {"compressors": [BloscCodec(cname=cname, clevel=clevel)]}
from numcodecs import Blosc # zarr 2
return {"compressor": Blosc(cname=cname, clevel=clevel)}
try:
import zarr
print(f"zarr {zarr.__version__}")
ZARR = True
except ImportError:
ZARR = False
print("zarr not installed — pip install zarr")
try:
import rechunker
print(f"rechunker {rechunker.__version__}")
RECHUNKER = True
except ImportError:
RECHUNKER = False
print("rechunker not installed — pip install rechunker")
WORK_DIR = Path(tempfile.mkdtemp(prefix="dask_recipe_tools_"))
print(f"Work dir: {WORK_DIR}")
zarr 3.4.0
rechunker not installed — pip install rechunker
Work dir: /tmp/dask_recipe_tools_m18mcink
1. Start Cluster
The dask_tmp path returned by setup_dask_client is on fast local storage ($PBS_JOBFS if
available). We will use it as the rechunker intermediate store.
client, cluster, dask_tmp = setup_dask_client(
workload_type="cpu",
max_workers=2,
reserve_mem_gb=2.0,
fallback_on_detection_failure=True,
dashboard=False,
)
print(f"Workers : {len(client.nthreads())}")
print(f"Temp dir : {dask_tmp}")
[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-3396
Workers: 2 | threads/worker: 1 | processes: True
Memory: total ~15.6 GiB | usable ~13.6 GiB | per-worker ~6.8 GiB
Compression: spill=auto | comm=False
Workers : 2
Temp dir : /tmp/dask-3396
2. Create a Sample NetCDF Dataset
Generate a modest time-series dataset that we will convert to Zarr.
import pandas as pd
N_TIME, N_LAT, N_LON = 500, 36, 72
np.random.seed(42)
temp = (288 + 8 * np.sin(2 * np.pi * np.arange(N_TIME) / 365)[:, None, None]
+ np.random.randn(N_TIME, N_LAT, N_LON)).astype("float32")
ds_source = xr.Dataset(
{"temperature": (["time", "lat", "lon"], temp,
{"units": "K", "long_name": "Temperature"})},
coords={
"time": pd.date_range("2010-01-01", periods=N_TIME, freq="D"),
"lat": (["lat"], np.linspace(-90, 90, N_LAT), {"units": "degrees_north"}),
"lon": (["lon"], np.linspace(-180, 180, N_LON), {"units": "degrees_east"}),
},
)
ds_source.attrs["conventions"] = "CF-1.8"
# Save as NetCDF
nc_path = WORK_DIR / "source.nc"
ds_source.to_netcdf(nc_path)
nc_mb = nc_path.stat().st_size / 1e6
print(f"Source NetCDF: {nc_path} ({nc_mb:.1f} MB)")
print(f"Dataset dims: {dict(ds_source.sizes)}")
Source NetCDF: /tmp/dask_recipe_tools_m18mcink/source.nc (5.2 MB)
Dataset dims: {'time': 500, 'lat': 36, 'lon': 72}
3. Direct NetCDF → Zarr Conversion (xarray)
The simplest route: open NetCDF, rechunk in memory, write Zarr. Works well for datasets that fit in the combined worker memory.
if ZARR:
zarr_out = WORK_DIR / "direct.zarr"
t0 = time.perf_counter()
ds_nc = xr.open_dataset(nc_path, chunks={"time": 100})
# Rechunk to analysis-friendly layout
ds_nc = ds_nc.chunk({"time": 250, "lat": N_LAT, "lon": N_LON})
# zarr.Blosc was removed in zarr 3; blosc_encoding() handles both versions
encoding = {"temperature": blosc_encoding()}
ds_nc.to_zarr(zarr_out, mode="w", consolidated=True, encoding=encoding)
ds_nc.close()
elapsed = time.perf_counter() - t0
zarr_mb = sum(f.stat().st_size for f in zarr_out.rglob("*") if f.is_file()) / 1e6
print(f"Direct conversion: {nc_mb:.1f} MB → {zarr_mb:.1f} MB ({elapsed:.2f}s)")
print(f"Compression ratio: {zarr_mb / nc_mb:.2f}x")
else:
print("zarr not available — skipping")
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/zarr/api/asynchronous.py:246: ZarrUserWarning: Consolidated metadata is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.
warnings.warn(
Direct conversion: 5.2 MB → 3.3 MB (0.67s)
Compression ratio: 0.63x
4. Memory-Bounded Rechunking with rechunk_dataset()
For datasets larger than available memory, rechunker streams data through a temporary store
under a strict memory budget (max_mem). dask_setup.rechunk_dataset() wraps it: the temporary
and output stores go to dask_tmp (fast local storage like $PBS_JOBFS), and when rechunker is
incompatible with the installed xarray or zarr (rechunker 0.5 predates zarr 3) it falls back to
xarray.to_zarr(), with a warning, since that path has no memory bound.
if ZARR and RECHUNKER:
rechunked_out = WORK_DIR / "rechunked.zarr"
# Source: 50 time steps per chunk. Target: time-contiguous chunks for time-series analysis.
ds_source_chunked = ds_source.chunk({"time": 50, "lat": N_LAT, "lon": N_LON})
t0 = time.perf_counter()
ds_rechunked = rechunk_dataset(
ds_source_chunked,
target_chunks={"time": 250, "lat": N_LAT, "lon": N_LON},
client=client,
dask_tmp=dask_tmp,
output_path=rechunked_out,
max_mem="100MB", # strict memory budget for rechunker
)
elapsed = time.perf_counter() - t0
out_mb = sum(f.stat().st_size for f in rechunked_out.rglob("*") if f.is_file()) / 1e6
print(f"Rechunked: {nc_mb:.1f} MB → {out_mb:.1f} MB ({elapsed:.2f}s)")
print(f"Output chunks: {dict(ds_rechunked.chunks)}")
elif not RECHUNKER:
print("rechunker not installed — pip install rechunker")
else:
print("zarr not installed — pip install zarr")
rechunker not installed — pip install rechunker
5. Validate Conversion Integrity
After any conversion, verify that the Zarr store matches the source data.
if ZARR and (WORK_DIR / "direct.zarr").exists():
zarr_out = WORK_DIR / "direct.zarr"
ds_src = xr.open_dataset(nc_path, chunks={})
ds_zarr = xr.open_zarr(zarr_out, consolidated=True)
# Compare statistics
src_mean = float(ds_src.temperature.mean().compute())
zarr_mean = float(ds_zarr.temperature.mean().compute())
src_std = float(ds_src.temperature.std().compute())
zarr_std = float(ds_zarr.temperature.std().compute())
print(f" {'NetCDF':>12} {'Zarr':>12} {'Match':>8}")
print("-" * 44)
print(f" mean {src_mean:>12.6f} {zarr_mean:>12.6f} {'✅' if abs(src_mean - zarr_mean) < 1e-3 else '❌'}")
print(f" std {src_std:>12.6f} {zarr_std:>12.6f} {'✅' if abs(src_std - zarr_std) < 1e-3 else '❌'}")
print(f" time steps {ds_src.sizes['time']:>12} {ds_zarr.sizes['time']:>12} "
f"{'✅' if ds_src.sizes['time'] == ds_zarr.sizes['time'] else '❌'}")
ds_src.close()
ds_zarr.close()
else:
print("zarr not available or conversion was skipped")
NetCDF Zarr Match
--------------------------------------------
mean 289.557617 289.557617 ✅
std 5.690793 5.690797 ✅
time steps 500 500 ✅
6. Incremental Time-Window Zarr Writes
For datasets too large to write at once, append time windows one at a time.
The first window uses mode="w" (creates the store); subsequent windows use append_dim="time".
if ZARR:
incremental_store = WORK_DIR / "incremental.zarr"
WINDOW = 100 # time steps per write
t0 = time.perf_counter()
for i, start in enumerate(range(0, N_TIME, WINDOW)):
stop = min(start + WINDOW, N_TIME)
chunk = ds_source.isel(time=slice(start, stop))
if i == 0:
chunk.to_zarr(incremental_store, mode="w", consolidated=False)
else:
chunk.to_zarr(incremental_store, append_dim="time")
# Consolidate metadata once at the end
zarr.consolidate_metadata(incremental_store)
elapsed = time.perf_counter() - t0
out_mb = sum(f.stat().st_size for f in incremental_store.rglob("*") if f.is_file()) / 1e6
# Verify
ds_check = xr.open_zarr(incremental_store, consolidated=True)
print(f"Incremental write: {out_mb:.1f} MB {elapsed:.2f}s")
print(f"Time steps in store: {ds_check.sizes['time']} (expected {N_TIME})")
ds_check.close()
else:
print("zarr not available — skipping")
Incremental write: 3.7 MB 0.11s
Time steps in store: 500 (expected 500)
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/zarr/api/asynchronous.py:246: ZarrUserWarning: Consolidated metadata is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.
warnings.warn(
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/zarr/api/asynchronous.py:246: ZarrUserWarning: Consolidated metadata is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.
warnings.warn(
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/zarr/api/asynchronous.py:246: ZarrUserWarning: Consolidated metadata is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.
warnings.warn(
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/zarr/api/asynchronous.py:246: ZarrUserWarning: Consolidated metadata is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.
warnings.warn(
client.close()
cluster.close()
shutil.rmtree(WORK_DIR, ignore_errors=True)
print("Cluster closed, temp files removed.")
Cluster closed, temp files removed.
Integration Best Practices
Rechunk through dask_tmp (fast local storage) with rechunk_dataset():
client, cluster, dask_tmp = setup_dask_client(workload_type="cpu")
ds_rechunked = rechunk_dataset(
ds,
target_chunks={"time": 250},
client=client,
dask_tmp=dask_tmp, # rechunker's temp store goes here
output_path="output.zarr",
max_mem="4GB", # at most ~50% of per-worker memory
)
Incremental writes for large time series:
for i, window in enumerate(time_windows):
if i == 0:
window.to_zarr("output.zarr", mode="w") # append_dim is only valid when appending
else:
window.to_zarr("output.zarr", append_dim="time")
zarr.consolidate_metadata("output.zarr") # once at the end
Zarr encoding recommendations:
lz4compressor for fast analysis readszstdfor archival (better compression ratio)Always
consolidated=Trueto avoid per-key metadata overhead
See the IO-Optimization wiki page for full details.