user-defined chunk grids #3750

Activity

1. TomNicholas commented on Mar 7, 2026

   

   TomNicholas

   on Mar 7, 2026

   Member

   More actions

   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).
2. maxrjones commented on Mar 9, 2026

   

   maxrjones

   on Mar 9, 2026

   Member

   More actions

   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.
3. dcherian commented on Mar 9, 2026

   

   dcherian

   on Mar 9, 2026

   Contributor

   More actions

   [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)

   👍React with 👍2Reacted by Tom Nicholas and Max Jones
4. LDeakin commented on Mar 9, 2026

   

   LDeakin

   on Mar 9, 2026

   Member

   More actions

   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.
5. dcherian commented on Mar 10, 2026

   

   dcherian

   on Mar 10, 2026

   Last edited by dcherian

   Contributor

   More actions

   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.
6. maxrjones commented on Mar 10, 2026

   

   maxrjones

   on Mar 10, 2026

   Member

   More actions

   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.
7. d-v-b commented on Mar 10, 2026

   

   d-v-b

   on Mar 10, 2026

   ContributorAuthor

   More actions

   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.
8. maxrjones commented on Mar 14, 2026

   

   maxrjones

   on Mar 14, 2026

   Member

   More actions

   I've been working on an alternative implementation for rectilinear chunk support, inspired by this thread and building on #3534. I posted details in #3534 (comment). This is the design doc, which also includes the implementation and links to integration prototypes for cubed, xarray, and virtualizarr. It addresses the tiled/periodic patterns discussed above (proven as simple to implement, but deferred pending a zarr extension definition).
9. maxrjones

   added a parent issue on Mar 20, 2026

   • ODD PI 26.2 Objective 5: 🤪 Expand virtualization support for quirky datasets NASA-IMPACT/veda-odd#308
10. maxrjones

    self-assigned this

    on Apr 6, 2026
11. maxrjones commented on Apr 7, 2026

    

    maxrjones

    on Apr 7, 2026

    Member

    More actions

    the outcomes from this discussion are detailed in https://github.com/zarr-developers/zarr-python/blob/main/design/chunk-grid.md
12. maxrjones

    closed this as completedon Apr 7, 2026