Optimizing Zarr Stores: Chunking, Compression, and Data Types#

Note

This example shows how to configure the encoding for Zarr stores created according to the v2 specification.

If you would prefer to use Zarr v3 stores, please refer to the Zarr v3 encoding tutorial.

When writing a Zarr store, you define an encoding, which is a set of rules that determine how the data is physically stored on disk. This includes how the data is divided into smaller pieces called “chunks”, how each chunk is compressed to reduce file size, and what precision is used to store numerical values. These choices directly affect read performance and storage efficiency without altering the actual data values.

Even if your Xarray dataset is well structured or your Dask array is already chunked, you still need to define encoding for Zarr. This is because Dask chunking and Zarr encoding serve different purposes. Dask chunks control how data is processed in memory during computation, while Zarr chunks define how data is laid out on disk for efficient reading and writing. Although these two concepts are independent, a properly encoded dataset allows for more efficient access with Dask as outlined in our Dask introduction.

The zarr-python library provides sensible default values for encoding parameters, which are usually sufficient for small datasets. However, when performance is critical, these defaults may no longer be optimal. In such cases, it is up to you, the data creator, to make deliberate decisions about chunking, compression, and data types. Taking control of the encoding allows you to tailor the storage format to your specific use case.

This article will introduce the concepts of chunking and compression, and discuss certain options for setting them for a dataset. As an example, we will store an IFS forecast in a Zarr store. The dataset includes both 2D and 3D fields, meaning we are working with multi-dimensional arrays along time, level, and cell dimensions.

import intake


cat = intake.open_catalog("https://tcodata.mpimet.mpg.de/internal.yaml")
ds = cat.HIFS(datetime="2024-09-01").to_dask()
ds

<xarray.Dataset> Size: 7GB
Dimensions:  (time: 64, cell: 196608, level: 13, crs: 1)
Coordinates:
  * time     (time) datetime64[ns] 512B 2024-09-01T03:00:00 ... 2024-09-11
  * level    (level) int64 104B 50 100 150 200 250 300 ... 600 700 850 925 1000
  * crs      (crs) float64 8B nan
Dimensions without coordinates: cell
Data variables: (12/39)
    100u     (time, cell) float32 50MB dask.array<chunksize=(6, 16384), meta=np.ndarray>
    100v     (time, cell) float32 50MB dask.array<chunksize=(6, 16384), meta=np.ndarray>
    10u      (time, cell) float32 50MB dask.array<chunksize=(6, 16384), meta=np.ndarray>
    10v      (time, cell) float32 50MB dask.array<chunksize=(6, 16384), meta=np.ndarray>
    2d       (time, cell) float32 50MB dask.array<chunksize=(6, 16384), meta=np.ndarray>
    2t       (time, cell) float32 50MB dask.array<chunksize=(6, 16384), meta=np.ndarray>
    ...       ...
    tp       (time, cell) float32 50MB dask.array<chunksize=(6, 16384), meta=np.ndarray>
    ttr      (time, cell) float32 50MB dask.array<chunksize=(6, 16384), meta=np.ndarray>
    u        (time, level, cell) float32 654MB dask.array<chunksize=(6, 1, 16384), meta=np.ndarray>
    v        (time, level, cell) float32 654MB dask.array<chunksize=(6, 1, 16384), meta=np.ndarray>
    vo       (time, level, cell) float32 654MB dask.array<chunksize=(6, 1, 16384), meta=np.ndarray>
    w        (time, level, cell) float32 654MB dask.array<chunksize=(6, 1, 16384), meta=np.ndarray>

Chunking#

We define multi-dimensional chunks for more efficient data access. The chunk size determines the number of elements after which an array will be cut into smaller chunks of data along a given dimension. Each chunk is then stored as a separate file within the Zarr store and hence can be loaded individually. A small horizontal chunk size, for example, allows you to access regional data without needing to load a whole global field.

The array shape and chunk size can be used to compute the amount of chunks (files) that will be created for an array:

\[ \left\lceil \frac{\text{Dimension size}}{\text{Chunk size}} \right\rceil = \text{Number of chunks} \]

For multi-dimensional arrays, the total amount of chunks is the product of the number of chunks in each dimension. The chunk size has significant impact:

  • Too small (e.g., 100 KB): Too many files – slow metadata handling, filesystem limits.

  • Too large (e.g., 1 GB): Loading a small region loads too much data – slow access.

We aim at a chunk size of about 8 MB which strikes a balance: fast to read, manageable file count.

def get_chunks(sizes):
    """Suggest a tuple of chunk sizes for a given dict of array dimensions."""
    match tuple(sizes.keys()):
        case ("time", "level", "cell"):
            chunks = {
                "time": 32,
                "level": 4,
                "cell": 4**7,
            }
        case ("time", "cell"):
            chunks = {
                "time": 32,
                "cell": 4**8,  # Account for missing `level` dimension
            }
        case (single_dim,):
            chunks = {
                single_dim: sizes[single_dim]
            }
        case _:
            chunks = {}

    return tuple((chunks[d] for d in sizes))


get_chunks(ds["tcwv"].sizes)

(32, 65536)

Warning

Depending on the total size of your dataset, a Zarr store may contain millions (!) of individual files, which might cause problems on some file systems.

Compression#

By default, data chunks will be compressed using the LZ4 compression algorithm. The LZ4 algorithm is chosen for its read and write speed.

Here, we compress all variables using Zstandard into a Blosc container instead. Blosc is a high-performance compression framework. It combines multiple compression algorithms with fast, parallelized filtering. Think of it as a “container” that holds compressed data efficiently. Zstandard offers a great balance between speed and compression ratio. Higher levels (like clevel=6) compress more aggressively, saving disk space with minimal slowdown.

import numcodecs


def get_compressor():
    """Return a numcodecs compressor."""
    return numcodecs.Blosc("zstd", clevel=6)


get_compressor()

Blosc(cname='zstd', clevel=6, shuffle=SHUFFLE, blocksize=0)

Data types#

Another setting to consider is the bit-precision of stored variables. Often times, arrays are stored in 64-bit floating point precision, which can result in large amounts of data on disk. Most meteorological variables (temperature, pressure, humidity) are well-represented in float32. In cases where this high accuracy is not needed, it is therefore beneficial to reduce the bit precision.

The following convenience function sets the output datatype to single-precision float for all given float subtypes.

import numpy as np


def get_dtype(da):
    """Suggest a dtype for storage (e.g., float32 for floats)."""
    if np.issubdtype(da.dtype, np.floating):
        # Store **all** floating point types as float32
        return "float32"
    else:
        # Preserve all other dtypes
        return da.dtype


get_dtype(ds["tcwv"])

'float32'

Changing the data type of your arrays to a lower precision (e.g. from float64 to float32) will change your data a bit (pun intended). When performing:

  • Long-term integrations (e.g., cumulative precipitation)

  • High-precision diagnostics (e.g., model diagnostics with small differences)

  • Scientific analysis requiring high accuracy

stick with float64 to avoid rounding errors.

Plug and play#

Finally, we can put the pieces together to define an encoding for the whole dataset. The following function loops over all variables and creates an encoding dictionary.

def get_encoding(dataset):
    return {
        var: {
            "compressor": get_compressor(),
            "dtype": get_dtype(dataset[var]),
            "chunks": get_chunks(dataset[var].sizes),
        }
        for var in dataset.variables
    }


get_encoding(ds[["t", "2t"]])

{'t': {'compressor': Blosc(cname='zstd', clevel=6, shuffle=SHUFFLE, blocksize=0),
  'dtype': 'float32',
  'chunks': (32, 4, 16384)},
 '2t': {'compressor': Blosc(cname='zstd', clevel=6, shuffle=SHUFFLE, blocksize=0),
  'dtype': 'float32',
  'chunks': (32, 65536)},
 'level': {'compressor': Blosc(cname='zstd', clevel=6, shuffle=SHUFFLE, blocksize=0),
  'dtype': dtype('<i8'),
  'chunks': (13,)},
 'time': {'compressor': Blosc(cname='zstd', clevel=6, shuffle=SHUFFLE, blocksize=0),
  'dtype': dtype('<M8[ns]'),
  'chunks': (64,)}}

The encoding dictionary can be passed to the to_zarr() method. When using Dask, make sure that the Dask chunks are aligned with the Zarr chunks (see Working with Dask). Otherwise the Zarr library will throw an error to prevent multiple Dask chunks from writing to the same chunk on disk.

# Prefer opening with an appropriate chunking over rechunking with `.chunk()`!
ds = cat.HIFS(
    datetime="2024-09-01",
    chunks={"time": 32, "cell": -1, "level": 4},
).to_dask()

ds.to_zarr(
    "test_dataset.zarr",
    encoding=get_encoding(ds),
    zarr_format=2,
    compute=False  # Set to True to store actual data
)

Delayed('_finalize_store-b9747ec2-0d1f-4050-b7da-7c0cf71aa662')

Tip

The compute=False option tells Dask to only write the Zarr metadata. This is convenient to verify that the encoding matches the dataset.