---
file_format: mystnb
kernelspec:
  name: python3
  display_name: Python 3
execution:
  timeout: 60
orphan: true
---

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

```{note}
This example shows how to configure the encoding for Zarr stores created according to the v3 specification.

If you would prefer to use Zarr v2 stores, please refer to [the Zarr v2 encoding tutorial](encoding-zarr-v2).
```

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](../dask).

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.

```{code-cell} ipython3
import intake


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

## Chunking

We define [multi-dimensional chunks](https://archive.unidata.ucar.edu/software/netcdf/workshops/most-recent/nc4chunking/WhatIsChunking.html) 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.

```{code-cell} ipython3
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)
```

```{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.
```

## Sharding

Zarr v3 introduces **sharding**, a feature that groups multiple data chunks into a single file (called a *shard*).
This reduces the total number of files (and inodes) on the filesystem, a major advantage for large datasets that would otherwise generate millions of small files.

The following function computes the **shard size in data elements** based on how many chunks you want to group.
Here, we write four `cell` and two `time` chunks into a single shard, effectively reducing the amount of files by a factor of eight.

```{code-cell} ipython3
def get_shards(sizes):
    """Suggest a tuple of shard sizes for a given dict of array dimensions."""
    chunk_sizes = get_chunks(sizes)

    # Map dimension names to chunk sizes
    dim_to_chunk = dict(zip(sizes.keys(), chunk_sizes))


    match tuple(sizes.keys()):
        case ("time", "level", "cell"):
            chunks = {
                "time": 2 * dim_to_chunk["time"],
                "level": dim_to_chunk["level"],
                "cell": 4 * dim_to_chunk["cell"],
            }
        case ("time", "cell"):
            chunks = {
                "time": 2 * dim_to_chunk["time"],
                "cell": 4 * dim_to_chunk["cell"],
            }
        case (single_dim,):
            chunks = {
                single_dim: sizes[single_dim]
            }
        case _:
            chunks = {}

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


get_shards(ds["tcwv"].sizes)
```

```{warning}
Although the Zarr v3 specification does not impose a hard limit on shard size, the `zarr-python` implementation may slow down significantly when a shard contains more than ~100 chunks.
```

## Compression

By default, data chunks will be compressed using the [LZ4](https://en.wikipedia.org/wiki/LZ4_(compression_algorithm)) compression algorithm.
The LZ4 algorithm is chosen for its read and write speed.

Here, we compress all variables using [Zstandard](https://en.wikipedia.org/wiki/Zstd) into a [Blosc](https://blosc.org) 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.

```{code-cell} ipython3
import zarr


def get_compressor():
    """Return a numcodecs compressor."""
    return zarr.codecs.BloscCodec(cname="zstd", clevel=6)


get_compressor()
```

## 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.

```{code-cell} ipython3
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"])
```

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.

```{code-cell} ipython3
def get_encoding(dataset):
    return {
        var: {
            "compressor": get_compressor(),
            "dtype": get_dtype(dataset[var]),
            "chunks": get_chunks(dataset[var].sizes),
            "shards": get_shards(dataset[var].sizes),
        }
        for var in dataset.variables
    }


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

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](../dask)).
Otherwise the Zarr library will throw an error to prevent multiple Dask chunks from writing to the same chunk on disk.

```{code-cell} ipython3
# 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=3,
    compute=False  # Set to True to store actual data
)
```
```{tip}
The `compute=False` option tells Dask to only write the Zarr **metadata**.
This is convenient to verify that the encoding matches the dataset.
```