Supporting Zarr V3 sharding with `.to_zarr` #9938

jacobbieker · 2025-01-10T13:01:59Z

jacobbieker
Jan 10, 2025

Hey,

Now that Zarr-Python V3 has been released, I was wondering if there is a way for Xarray to write sharded Zarrs at the moment? It would be quite helpful for some of the workflows I am working with.

So far, I've somewhat gotten an example working with a few changes to xarray, passing shards as part of the variable encoding. The Zarr starts getting written, but then has a checksum error.

Answered by jhamman

Jan 13, 2025

Thanks for the explanation. You are right out there on the bleeding edge using DataTree + v3-sharding! Flagging two potential issues:

The Zarr3+DataTree interface may not be fully operational. This interface is probably the least mature part of the stack. It does seem like its working for you though which is great.
Xarray may need to update how it writes sharded arrays when using Dask. Currently, you must write full shards, this means that if you are using Dask, your Dask-chunks need to correspond to Zarr-shards. I suspect this is your current issue.

View full answer

jhamman · 2025-01-10T14:13:25Z

jhamman
Jan 10, 2025
Maintainer

Hi @jacobbieker -- great question! We should update the Xarray docs that describe using Zarr for how to do this. Do mind sharing the error you ran into?

3 replies

jacobbieker Jan 10, 2025
Author

Yes, the first error I got was from here:

xarray/xarray/backends/zarr.py

Line 449 in 5279bd1

valid_encodings = {

where having a shards entry to pass to the backend raises an error. Adding shards there allowed it to continue.

The error after that was on the writing.

Here is the full traceback:

Traceback (most recent call last):
  File "/Users/jacob/Development/planetary-datasets/planetary_datasets/provider/eumetsat/mtg.py", line 176, in <module>
    download_and_process_normal_and_high_res(start, end)
  File "/Users/jacob/Development/planetary-datasets/planetary_datasets/provider/eumetsat/mtg.py", line 72, in download_and_process_normal_and_high_res
    save_datatree(mtg_datatree["/hr_05km"], mtg_datatree["/hr_1km"], mtg_datatree["/1km"], mtg_datatree["/2km"])
  File "/Users/jacob/Development/planetary-datasets/planetary_datasets/provider/eumetsat/mtg.py", line 173, in save_datatree
    full_dataset.to_zarr(f"mtg_datatree_v3_{timestamp}.zarr", mode="w", compute=True, encoding=encoding, zarr_format=3)
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/xarray/core/datatree.py", line 1699, in to_zarr
    _datatree_to_zarr(
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/xarray/core/datatree_io.py", line 121, in _datatree_to_zarr
    ds.to_zarr(
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/xarray/core/dataset.py", line 2622, in to_zarr
    return to_zarr(  # type: ignore[call-overload,misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/xarray/backends/api.py", line 2217, in to_zarr
    writes = writer.sync(
             ^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/xarray/backends/common.py", line 358, in sync
    delayed_store = chunkmanager.store(
                    ^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/xarray/namedarray/daskmanager.py", line 247, in store
    return store(
           ^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/array/core.py", line 1282, in store
    compute_as_if_collection(Array, store_dsk, map_keys, **kwargs)
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/base.py", line 397, in compute_as_if_collection
    return schedule(dsk2, keys, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/threaded.py", line 91, in get
    results = get_async(
              ^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/local.py", line 516, in get_async
    raise_exception(exc, tb)
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/local.py", line 324, in reraise
    raise exc
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/local.py", line 229, in execute_task
    result = task(data)
             ^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/_task_spec.py", line 651, in __call__
    return self.func(*new_argspec)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/array/core.py", line 4575, in store_chunk
    return load_store_chunk(x, out, index, lock, return_stored, False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/dask/array/core.py", line 4557, in load_store_chunk
    out[index] = x
    ~~~^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/array.py", line 2464, in __setitem__
    self.set_orthogonal_selection(pure_selection, value, fields=fields)
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/_compat.py", line 43, in inner_f
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/array.py", line 2920, in set_orthogonal_selection
    return sync(
           ^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/sync.py", line 142, in sync
    raise return_result
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/sync.py", line 98, in _runner
    return await coro
           ^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/array.py", line 1354, in _set_selection
    await self.codec_pipeline.write(
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/codec_pipeline.py", line 468, in write
    await concurrent_map(
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/common.py", line 68, in concurrent_map
    return await asyncio.gather(*[asyncio.ensure_future(run(item)) for item in items])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/common.py", line 66, in run
    return await func(*item)
           ^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/codec_pipeline.py", line 344, in write_batch
    await self.encode_partial_batch(
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/codec_pipeline.py", line 229, in encode_partial_batch
    await self.array_bytes_codec.encode_partial(batch_info)
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/abc/codec.py", line 235, in encode_partial
    await concurrent_map(
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/common.py", line 68, in concurrent_map
    return await asyncio.gather(*[asyncio.ensure_future(run(item)) for item in items])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/common.py", line 66, in run
    return await func(*item)
           ^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/codecs/sharding.py", line 582, in _encode_partial_single
    await self._load_full_shard_maybe(
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/codecs/sharding.py", line 730, in _load_full_shard_maybe
    await _ShardReader.from_bytes(shard_bytes, self, chunks_per_shard)
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/codecs/sharding.py", line 194, in from_bytes
    obj.index = await codec._decode_shard_index(shard_index_bytes, chunks_per_shard)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/codecs/sharding.py", line 632, in _decode_shard_index
    await get_pipeline_class()
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/codec_pipeline.py", line 435, in decode
    output.extend(await self.decode_batch(batch_info))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/codec_pipeline.py", line 172, in decode_batch
    chunk_bytes_batch = await bb_codec.decode(
                        ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/abc/codec.py", line 129, in decode
    return await _batching_helper(self._decode_single, chunks_and_specs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/abc/codec.py", line 407, in _batching_helper
    return await concurrent_map(
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/common.py", line 68, in concurrent_map
    return await asyncio.gather(*[asyncio.ensure_future(run(item)) for item in items])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/core/common.py", line 66, in run
    return await func(*item)
           ^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/abc/codec.py", line 420, in wrap
    return await func(chunk, chunk_spec)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jacob/.conda/envs/pc/lib/python3.12/site-packages/zarr/codecs/crc32c_.py", line 46, in _decode_single
    raise ValueError(
ValueError: Stored and computed checksum do not match. Stored: b'\x11\x1c\xd5\x03'. Computed: b'\xb4P\t\x94'.

For reference, the code I'm using to try to write the Zarr is here: https://github.com/jacobbieker/planetary-datasets/blob/a87b70d639355e19e3474b4509373560f206050f/planetary_datasets/provider/eumetsat/mtg.py#L150-L173

I am using a DataTree, which does write a non-sharded Zarr V3 correctly.

This is the DataTree:

<xarray.DataTree>
Group: /
├── Group: /hr_1km
│       Dimensions:    (time: 1, y: 11136, x: 11136)
│       Coordinates:
│         * time       (time) datetime64[ns] 8B 2025-01-02T00:05:00
│         * y          (y) float64 89kB -5.567e+06 -5.566e+06 ... 5.566e+06 5.567e+06
│         * x          (x) float64 89kB -5.567e+06 -5.566e+06 ... 5.566e+06 5.567e+06
│           longitude  (y, x) float64 992MB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
│           latitude   (y, x) float64 992MB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
│       Data variables:
│           ir_105     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           ir_38      (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│       Attributes:
│           history:      Created by pytroll/satpy on 2025-01-10 15:07:04.655705
│           Conventions:  CF-1.7
├── Group: /hr_05km
│       Dimensions:    (time: 1, y: 22272, x: 22272)
│       Coordinates:
│         * time       (time) datetime64[ns] 8B 2025-01-02T00:05:00
│         * y          (y) float64 178kB -5.568e+06 -5.567e+06 ... 5.567e+06 5.568e+06
│         * x          (x) float64 178kB -5.568e+06 -5.567e+06 ... 5.567e+06 5.568e+06
│           longitude  (y, x) float64 4GB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
│           latitude   (y, x) float64 4GB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
│       Data variables:
│           nir_22     (time, y, x) float32 2GB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           vis_06     (time, y, x) float32 2GB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│       Attributes:
│           history:      Created by pytroll/satpy on 2025-01-10 15:06:53.192232
│           Conventions:  CF-1.7
├── Group: /1km
│       Dimensions:    (time: 1, y: 11136, x: 11136)
│       Coordinates:
│         * time       (time) datetime64[ns] 8B 2025-01-02T00:05:00
│         * y          (y) float64 89kB -5.567e+06 -5.566e+06 ... 5.566e+06 5.567e+06
│         * x          (x) float64 89kB -5.567e+06 -5.566e+06 ... 5.566e+06 5.567e+06
│           longitude  (y, x) float64 992MB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
│           latitude   (y, x) float64 992MB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
│       Data variables:
│           nir_13     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           nir_16     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           nir_22     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           vis_04     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           vis_05     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           vis_06     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           vis_08     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│           vis_09     (time, y, x) float32 496MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
│       Attributes:
│           history:      Created by pytroll/satpy on 2025-01-10 15:06:41.743445
│           Conventions:  CF-1.7
└── Group: /2km
        Dimensions:    (time: 1, y: 5568, x: 5568)
        Coordinates:
          * time       (time) datetime64[ns] 8B 2025-01-02T00:05:00
          * y          (y) float64 45kB -5.567e+06 -5.565e+06 ... 5.565e+06 5.567e+06
          * x          (x) float64 45kB -5.567e+06 -5.565e+06 ... 5.565e+06 5.567e+06
            longitude  (y, x) float64 248MB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
            latitude   (y, x) float64 248MB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
        Data variables:
            ir_105     (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
            ir_123     (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
            ir_133     (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
            ir_38      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
            ir_87      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
            ir_97      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
            wv_63      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
            wv_73      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
        Attributes:
            history:      Created by pytroll/satpy on 2025-01-10 15:06:47.895700
            Conventions:  CF-1.7

jhamman Jan 13, 2025
Maintainer

Thanks for the explanation. You are right out there on the bleeding edge using DataTree + v3-sharding! Flagging two potential issues:

The Zarr3+DataTree interface may not be fully operational. This interface is probably the least mature part of the stack. It does seem like its working for you though which is great.
Xarray may need to update how it writes sharded arrays when using Dask. Currently, you must write full shards, this means that if you are using Dask, your Dask-chunks need to correspond to Zarr-shards. I suspect this is your current issue.

Answer selected by jacobbieker

jacobbieker Jan 13, 2025
Author

Okay, thanks! I've now been able to get a single Dataset write a sharded V3 Zarr to disk, writing one of the Datasets in the DataTree on its own works at least. I just needed the Zarr-shard to be evenly divisible by the Dask-chunk. So it doesn't seem like the Dask-chunk has to be the same size as the Zarr-shard.

Here is the output from writing one of the Datasets in the above DataTree on its own as a sharded Zarr V3. Each data variable is sharded as (1,5568,5568), while the Dask-chunks are (1,1392,1392):

<xarray.Dataset> Size: 1GB
Dimensions:    (time: 1, y: 5568, x: 5568)
Coordinates:
  * time       (time) datetime64[ns] 8B 2025-01-01T11:55:00
    longitude  (y, x) float64 248MB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
  * x          (x) float64 45kB -5.567e+06 -5.565e+06 ... 5.565e+06 5.567e+06
    latitude   (y, x) float64 248MB dask.array<chunksize=(1392, 1392), meta=np.ndarray>
  * y          (y) float64 45kB -5.567e+06 -5.565e+06 ... 5.565e+06 5.567e+06
Data variables:
    ir_133     (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
    ir_97      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
    ir_105     (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
    wv_73      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
    ir_38      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
    wv_63      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
    ir_87      (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
    ir_123     (time, y, x) float32 124MB dask.array<chunksize=(1, 1392, 1392), meta=np.ndarray>
Attributes:
    history:      Created by pytroll/satpy on 2025-01-13 16:32:55.674674
    Conventions:  CF-1.7

For now, I can just write each Dataset independently, the DataTree was more for ease of use with these. Thanks for your help!

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Supporting Zarr V3 sharding with `.to_zarr` #9938

{{title}}

Replies: 1 comment 3 replies

{{title}}

{{title}}

{{title}}

{{title}}

Select a reply

Supporting Zarr V3 sharding with .to_zarr #9938

jacobbieker Jan 10, 2025

Replies: 1 comment · 3 replies

jhamman Jan 10, 2025 Maintainer

jacobbieker Jan 10, 2025 Author

jhamman Jan 13, 2025 Maintainer

jacobbieker Jan 13, 2025 Author

Supporting Zarr V3 sharding with `.to_zarr` #9938

jacobbieker
Jan 10, 2025

Replies: 1 comment 3 replies

jhamman
Jan 10, 2025
Maintainer

jacobbieker Jan 10, 2025
Author

jhamman Jan 13, 2025
Maintainer

jacobbieker Jan 13, 2025
Author