Modifying sharded datasets causes Dask error

[taus]

In an effort to rechunk and shard some datasets i am first rechunking the dataset. Later i want reopen the dataset to add some metadata to the dataset. When doing this i get an error that my chunks would overlap Dask chunks

I am chunking the dataset by setting the encoding as that is the only way i have found to store dataset with shards.
First i set the dataset dask chunks to the size of the shard. Then i am setting the chunks and shards in the zarr encoding and finally i am writing to the zarr store. This works just fine and the store is saved as expected.

Next i try to open the dataset, modify some metadata and write it back, causing the error.

I have created a minimal example illustrating the issue

import numpy as np
import xarray as xr

store = "test.zarr"

chunks= {"x":2,"y":2}
shards= {"x":4,"y":4}
dataset_shape = {"x":10,"y":10}
try:
    rng = np.random.default_rng()
    ds = xr.Dataset(
        {
            "data": (("x", "y"),rng.random((dataset_shape["x"], dataset_shape["y"]))),
        },
        coords={
            "x": np.arange(dataset_shape["x"]),
            "y": np.arange(dataset_shape["y"]),
        }
    )

    # Make the xarray chunks match the zarr shards
    ds_sharded = ds.chunk(shards)

    encoding = {
        "data": {
            "chunks": (2, 2),
            "shards": (4, 4)
        },
        "x": {
            "chunks": (2,),
            "shards": (4,)
        },
        "y": {
            "chunks": (2,),
            "shards": (4,)
        },
    }

    try:
        ds_sharded.to_zarr(
            store,
            mode="w",
            encoding=encoding,
            zarr_format=3,
            consolidated=False,
        )
        print("Zarr store created without errors.")
    except ValueError as e:
        print(f"Error: {e}")

    #reopen dataset and store it back
    reopend_ds = xr.open_zarr(store,consolidated=False)
    print(f"ds_sharded.chunks: {ds_sharded.chunks}")
    print(f"ds2.chunks: {reopend_ds.chunks}")
    #This throws an error
    reopend_ds.to_zarr(store, mode="a")

except Exception as e:
    print(f"Error: {e}")

Running this i get

Zarr store created without errors.
ds_sharded.chunks: Frozen({'x': (4, 4, 2), 'y': (4, 4, 2)})
ds2.chunks: Frozen({'x': (2, 2, 2, 2, 2), 'y': (2, 2, 2, 2, 2)})
Error: Specified Zarr chunks encoding['chunks']=(4, 4) for variable named 'data' would overlap multiple Dask chunks. Please check the Dask chunks at position 1 and 2, on axis 0, they are overlapped on the same Zarr chunk in the region slice(None, None, None). Writing this array in parallel with Dask could lead to corrupted data. To resolve this issue, consider one of the following options: - Rechunk the array using `chunk()`. - Modify or delete `encoding['chunks']`. - Set `safe_chunks=False`. - Enable automatic chunks alignment with `align_chunks=True`.

The issue seems to be that when reading back the dataset, the Dask chunks are aligned with the zarr chunks instead of the Zarr shards.

Is this expected behavior? It seems like reopening the dataset, xarray looses information about the sharding.

[ianhi]

Is this expected behavior? It seems like reopening the dataset, xarray looses information about the sharding.

I would say no, we should not be writing data we cannot read back.

I thin kthis is likely at least related to this isseu: #10840

[hexsyro]

The chunk mismatch on reopen is expected — xarray reads back zarr chunks, not shards. The shard metadata is preserved in zarr but xarray doesn't map it back to Dask chunks on open. Options that avoid the error without disabling safety:

ds = xr.open_zarr(store, chunks=shards) — force chunk alignment on open
reopend_ds.chunk(shards).to_zarr(store, mode="a") — rechunk before writing
Set encoding again on write to match original sharding
Option 1 is cleanest if you know your shard sizes upfront. Option 3 is the pattern we use internally at Social Intel for managing large sharded datasets through our enrichment pipeline.

[khlaifiabilel]

For the stated second step, if “add some metadata” means dataset or variable attributes, there is no need to send the lazily opened arrays back through Dataset.to_zarr at all. Update only the Zarr metadata:

import zarr

root = zarr.open_group(
    store,
    mode="a",
    zarr_format=3,
    use_consolidated=False,
)

# xr.Dataset.attrs
root.attrs.update({"title": "example", "history": "metadata added"})

# xr.DataArray.attrs for a stored variable, if needed
root["data"].attrs.update({"units": "1", "long_name": "example data"})

A subsequent xr.open_zarr(store, consolidated=False) reads those attributes. This does not schedule any array-data writes, so Dask/Zarr chunk safety validation is not involved and the existing chunks and shards remain unchanged. Attribute values must be JSON-serializable.

This is appropriate for user/CF-style attributes only. Structural changes such as dimensions, shapes, chunking, sharding, codecs, coordinates, or encoded dtype are not attribute-only changes and should still go through an aligned data-writing path such as reopening/rechunking with shard-aligned Dask chunks. If consolidated metadata is enabled for another store, it must be refreshed after direct metadata changes; the example store was created with consolidated=False.