Question about loadable_variables for leap-year daily-data concatenation

[simentha]

Hey! I was very interested reading about this package as I have always wanted a way to "cache" open_mfdataset from xarray. My use case is that I have big datacubes for climate data sharing the exact same grid and data_vars, but they are divided into yearly zarr stores. As you probably can tell, I get an issue when I try to concatenate these stores into a virtual dataset because the chunk shape is inconsistent in leap-years (8784 vs 8760). In the FAQ it is written that:

Some of your variables have inconsistent-length chunks, and you want to be able to concatenate them together. For example you might have multiple virtual datasets with coordinates of inconsistent length (e.g., leap years within multi-year daily data). Loading them allows you to rechunk them however you like.

but I could not really get this to work. I am trying to create a minimal example - and would love to contribute with this as example documentation in a jupyter notebook if you'd like - but I am not sure if it is possible to do what I want here.

I have created a mini example to reproduce a similar use case for what I want. First, create couple of dummy HDF datasets with different time-axis length:

from datetime import datetime, timedelta
import numpy as np

january_times = [datetime(2025, 1, 1) + timedelta(hours=i) for i in range(0, 744)]
february_times = [datetime(2025, 2, 1) + timedelta(hours=i) for i in range(0, 672)]
lat = np.linspace(0, 10, 10)
lon = np.linspace(0, 10, 10)

ds_january = xr.Dataset(
    coords=dict(time=january_times, lon=lon, lat=lat),
    data_vars=dict(temperature=(("time", "lon", "lat"), np.ones((744, 10, 10)))),
).chunk(dict(time=-1)).to_netcdf("january.nc")

ds_february = xr.Dataset(
    coords=dict(time=february_times, lon=lon, lat=lat),
    data_vars=dict(temperature=(("time", "lon", "lat"), np.ones((672, 10, 10)))),
).chunk(dict(time=-1)).to_netcdf("february.nc")

Then try to concat them using VirtualiZarr and xarray:

from obstore.store import LocalStore

from virtualizarr import open_virtual_dataset
from virtualizarr.parsers import HDFParser
from virtualizarr.registry import ObjectStoreRegistry
import xarray as xr

from pathlib import Path

store_path = Path.cwd()
file_paths = [str(store_path / "january.nc"), str(store_path / "february.nc")]
file_urls = [f"file://{file_path}" for file_path in file_paths]

store = LocalStore(prefix=store_path)
registry = ObjectStoreRegistry({file_url: store for file_url in file_urls})
parser = HDFParser()
datasets = []
for url in file_urls:
    vds = open_virtual_dataset(
        url=url, parser=parser, registry=registry, loadable_variables=["time", "lat", "lon"]
    )
    datasets.append(vds)

mds = xr.concat(datasets, dim="time")
mds

This raises the exception:

ValueError: Cannot concatenate arrays with inconsistent chunk shapes: (672, 10, 10) vs (744, 10, 10) .Requires ZEP003 (Variable-length Chunks).

I have tried to re-chunk the vds before concat, but not yet figured out how this should be done. Is what I try to do here possible?

[sharkinsspatial]

@simentha Thank you for concise example. A few notes,

1. As the exception stated this is an upstream issue with the Zarr specification that is outlined in https://zarr.dev/zeps/draft/ZEP0003.html. Luckily there is a Zarr Developer Summit this week where this spec limitation is being addressed with a new Zarr extension Rectilinear (variable-length) chunk grid zarr-extensions#25. Once this lands in zarr-python we'll be able to support concatenating virtual arrays with inconsistent chunk shapes.
2. In the interim, there is not a good solution for purely virtual arrays. You mentioned re-chunking the virtual dataset. By its nature, because we are referencing existing chunks in archival files, a virtual dataset itself cannot be re-chunked. Instead, you would need to load all of the chunks, re-chunk them and store them as a new Zarr dataset with native rather than virtual chunks. We have some updates coming to the documentation and we'll try to to clarify this point more since we've seen confusion about this a few times.

[simentha]

@sharkinsspatial Thank you so much for the clarification! My understanding about the zarr format was that it currently should not work - but I misunderstood the docs and got unsure if there was a "magic" workaround for this in the virtual layer.

I think it is not too much work to rewrite the zarr store to one giant cube with unified chunks and write using xarray.to_zarr with the "region" keyword 👍 The reason for the yearly structure was to have a "clean" 1:1 mapping between archive tier data (NetCDF) and cloud-optimized Zarr store and it felt a bit safer for parallel write processes.