user-defined chunk grids

[d-v-b]

In the v3 spec, the chunk_grid field is extensible. People can define their own chunk grids.

In Zarr python, we have 1 chunk grid right (the regular chunk grid, every chunk is the same shape), and we are planning on adding another one (the rectilinear chunk grid, every chunk can have a different shape). Over in Zarrs, there's a third chunk grid, which is regular except for boundary chunks, which are trimmed. We should support this too.

The rectilinear chunk grid is strictly more general than both the regular chunk grid and the zarrs regular-bounded chunk grid -- any regular or regular-bounded chunk grid can be expressed as a rectilinear chunk grid, and if your implementation supports the rectilinear chunk grid, you support these other two chunk grids for free.

This is very different from other extension points like data types or chunk key encodings. Supporting uint8 doesn't help you with float32 or vice versa.

Over in #3735 I set up chunk grids the same way as the other extensible metadata, with a registry. This was basically extrapolation from how data types and chunk key encodings work. But @maxrjones was skeptical about this approach, and I think that skepticism was right. I don't think we can re-use the chunk key encoding / data patterns for chunk grids.

Here are some thoughts about how our chunk grid API could look when we add support for rectilinear chunking. I'm interested in any and all feedback on these ideas:

• we should have all our chunk grid logic in one place, e.g. just one chunk grid class. The regular chunk grid doesn't need a separate class, because it's just a special case of rectilinear chunking (rectilinear chunk grids can also be regular). We should think about the rectilinear chunk grid implementation as a replacement for our current (regular) chunk grid implementation.
• Unless directed otherwise we should serialize array metadata with the simplest chunk grid metadata that can describe the user-requested chunk arrangement of the array. If a user asks for regular chunks, we serialize the regular chunk grid. if a user has a regularly-chunked array and resizes it, we should expose control over the chunking of the expanded / contracted region, and update the chunk grid metadata accordingly. It should be easy, but also transparent, to opt into rectilinear chunking starting from a regularly-chunked array.
• Because the rectilinear chunk grid will be the most general chunk grid we can support as an implementation, and other known chunk grids are special cases of the rectilinear chunk grid, I don't think there's value in allowing user-defined chunk grids. Happy to be corrected here if there's some unknown chunk grid that's more general than the rectilinear chunk grid which we could actually support. I'm not aware of one that retains the properties we assume for chunk grids.

Anyway those are some of my thoughts while thinking about how to ship the features added in #3534, curious to hear if people have other ideas

[TomNicholas]

Over in Zarrs, there's a third chunk grid, which is regular except for boundary chunks, which are trimmed. We should support this too.

This is very interesting because one of the motivating cases for variable-length chunk grids in VirtualiZarr is very similar to this. HPC fluid simulation codes often have grids with uniform chunks shape except at the boundary, where the chunks are larger, to accomodate setting boundary conditions. So very similar but not quite the same to the zarrs case. This comes up when parallelizing, e.g. across 5 nodes on HPC, and the chunk lengths along one dimension might look like [10, 8, 8, 8, 10]. Oceanographic models like ROMS have this structure (see zarr-developers/VirtualiZarr#217).

[maxrjones]

Great thoughts @d-v-b! I wrote out a prospectus based on this idea to better understand how it interacts with the user-facing API, indexing, and sharding, along with a proposed implementation path. I particularly like how it might simplify variable-length chunks within shards.

[dcherian]

[10, 8, 8, 8, 10].

This could also be handled by the run length encoding in the MVP.

A pattern that is missing in the MVP is np.tile (RLE is np.repeat). This could be useful for periodic or quasi-periodic chunk patterns found along temporal axes. E.g. days in a month (with or without leap years)

[LDeakin]

This could be useful for periodic or quasi-periodic chunk patterns found along temporal axes

Temporal is a complex pattern to shoehorn into something like this. What about a temporal chunk grid that is date/time aware? It could layer over another chunk grid for non-temporal axes.

[dcherian]

What about a temporal chunk grid that is date/time aware?

I was skeptical at first but agree that this would be a good use-case for an external chunk grid definition. In particular, becasue in the Earth Sciences, we use many custom calendars. This could play well with Xarray's ability to rechunk to a frequency.

That said, I think Zarr-python should support an encoding that is basically np.tile, it's a generalization of the regular chunk grid.

[maxrjones]

I put together a quick prototype of a TiledDimension in maxrjones@9c0f582 and pydata/xarray@3e27417.

UV script demonstration of periodic chunk patterns

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "zarr @ git+https://github.com/maxrjones/zarr-python@poc/unified-chunk-grid",
#     "xarray @ git+https://github.com/maxrjones/xarray@poc/unified-zarr-chunk-grid",
#     "dask[array]",
# ]
# ///
"""
Demo: TiledDimension for periodic/quasi-periodic chunk patterns.

RLE (np.repeat) compresses runs of identical values:
    [10, 10, 10, 20, 20] -> [[10, 3], [20, 2]]

Tile (np.tile) compresses repeating *patterns*:
    [31, 28, 31, ...] x 30 -> {"tile": [31, 28, ..., 31], "repeat": 30}

This is useful for temporal axes where chunk sizes follow calendar patterns
(days per month, hours per day, etc).
"""

import json
import tempfile
from pathlib import Path

import dask.array as da
import numpy as np
import xarray as xr
import zarr
from zarr.core.chunk_grids import TiledDimension, serialize_chunk_grid

# -- Calendar constants -------------------------------------------------------
MONTHS_365 = (31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31)
N_YEARS = 30
N_DAYS = 365 * N_YEARS  # 10950

# Dask chunks following the monthly pattern: each chunk = one month
time_chunks = MONTHS_365 * N_YEARS  # 360 chunk sizes

with tempfile.TemporaryDirectory() as tmp:
    store = Path(tmp) / "demo.zarr"

    # -- 1. Create dataset with monthly dask chunks ----------------------------
    time = xr.date_range("2000-01-01", periods=N_DAYS, freq="D")
    temperature = xr.DataArray(
        da.random.random((N_DAYS, 180, 360), chunks=(time_chunks, 180, 360)),
        dims=["time", "lat", "lon"],
        coords={"time": time},
    )
    ds = xr.Dataset({"temperature": temperature})
    print(f"Dataset: {dict(ds.sizes)}")
    print(f"Dask chunks along time: {len(time_chunks)} chunks (one per month)")

    # -- 2. Write to zarr v3 (preserves variable chunks) -----------------------
    ds.to_zarr(store, zarr_format=3, mode="w", consolidated=False)

    # -- 3. Inspect the on-disk chunk grid -------------------------------------
    z = zarr.open_array(store / "temperature")
    grid = z.metadata.chunk_grid
    serialized = serialize_chunk_grid(grid, z.metadata.chunk_grid_name)
    chunk_shapes = serialized["configuration"]["chunk_shapes"]

    print(f"\nOn-disk chunk grid ({z.metadata.chunk_grid_name}):")
    time_spec = chunk_shapes[0]
    if isinstance(time_spec, dict):
        print(f"  time: tile encoding -- pattern of {len(time_spec['tile'])} edges, "
              f"repeated {time_spec['repeat']}x")
        print(f"  Serialized: {json.dumps(time_spec)}")
    else:
        print(f"  time: {type(time_spec).__name__} with {len(time_spec)} entries")
    print(f"  lat:  {chunk_shapes[1]}")
    print(f"  lon:  {chunk_shapes[2]}")

    # -- 4. Read back and verify round-trip ------------------------------------
    ds2 = xr.open_zarr(store, consolidated=False)
    print(f"\nRound-trip chunks along time: {len(ds2.temperature.chunks[0])} chunks")
    assert ds2.temperature.chunks is not None
    assert ds2.temperature.chunks[0] == time_chunks
    print("  Monthly chunk pattern preserved")

    # -- 5. Indexing: slice across chunk boundaries ----------------------------
    # Select Feb-Apr 2015 (crosses 3 monthly chunk boundaries)
    sel = ds2.temperature.sel(time=slice("2015-02-01", "2015-04-30"))
    result = sel.values
    print(f"\nsel(time='2015-02-01':'2015-04-30'): shape={sel.shape}, "
          f"{result.nbytes / 1024:.0f} KB")

    # Fancy indexing: first day of each month for 5 years
    first_days = [
        np.datetime64(f"{y}-{m:02d}-01")
        for y in range(2000, 2005) for m in range(1, 13)
    ]
    sel2 = ds2.temperature.sel(time=first_days)
    result2 = sel2.values
    print(f"sel(first-of-month, 5 years):     shape={sel2.shape}, "
          f"{result2.nbytes / 1024:.0f} KB")

    # -- 6. TiledDimension internals -------------------------------------------
    dim = grid.dimensions[0]
    if isinstance(dim, TiledDimension):
        full_json = json.dumps(serialized, separators=(",", ":"))
        flat = list(dim.edges)
        flat_json = json.dumps({
            "name": "rectilinear",
            "configuration": {"chunk_shapes": [flat, chunk_shapes[1], chunk_shapes[2]]},
        }, separators=(",", ":"))
        print(f"\nTiledDimension: pattern={len(dim.pattern)} edges, "
              f"repeats={dim.repeats}, total={dim.nchunks} chunks")
        print(f"  In-memory:  {len(dim.pattern)} edges (vs {dim.nchunks} for VaryingDimension)")
        print(f"  On-disk:    {len(full_json)} bytes (vs {len(flat_json)} flat) "
              f"= {len(flat_json) / len(full_json):.1f}x smaller")

A key component of https://github.com/maxrjones/zarr-python/tree/poc/unified-chunk-grid is that the chunk grid is broken down by dimension, which really helps here because you're likely to have a repeating pattern over one but not all dimensions.

It doesn't really answer the question of whether the chunk grids should be pluggable though. My intuition is to make sure the initial release doesn't preclude a future pluggable chunk grid feature, but to not include it right away. I also wouldn't want to include TiledDimensions in the initial release despite the working prototype. I would rather a decision is made about the serialization format first. Still, it's useful to know that the per-dimension design in https://hackmd.io/B72LaHyNR8Co0vq9VXSFHA?view supports this without changes to indexing or the codec pipeline.

[d-v-b]

That said, I think Zarr-python should support an encoding that is basically np.tile, it's a generalization of the regular chunk grid.

Do you mean as user-facing API, or something in the metadata? I don't think we want to innovate on the metadata side in zarr python, that should go through a change to the proposed spec. And without changes to the metadata, a tiled chunk grid would not have an efficient encoding in metadata.