DataTree.to_zarr() is very slow writing to high latency store

What is your issue?

Repost of xarray-contrib/datatree#277, with some updates.

Test case

Write a tree containing 13 nodes and negligible data to S3/GCS with fsspec:

import numpy as np
import xarray as xr

ds = xr.Dataset(
    data_vars={
        "a": xr.DataArray(np.ones((2, 2)), coords={"x": [1, 2], "y": [1, 2]}),
        "b": xr.DataArray(np.ones((2, 2)), coords={"x": [1, 2], "y": [1, 2]}),
        "c": xr.DataArray(np.ones((2, 2)), coords={"x": [1, 2], "y": [1, 2]}),
    }
)

dt = xr.core.datatree.DataTree()
for first_level in [1, 2, 3]:
    dt[f"{first_level}"] = DataTree(ds)
    for second_level in [1, 2, 3]:
        dt[f"{first_level}/{second_level}"] = DataTree(ds)

%time dt.to_zarr("test.zarr", mode="w")

bucket = "s3|gs://your-bucket/path" 
%time dt.to_zarr(f"{bucket}/test.zarr", mode="w")

Gives:

CPU times: user 287 ms, sys: 43.9 ms, total: 331 ms
Wall time: 331 ms
CPU times: user 3.22 s, sys: 219 ms, total: 3.44 s
Wall time: 1min 4s

This is a bit better than in the original issue due to improvements elsewhere in the stack, but still really slow for heavily nested but otherwise small datasets.

Potential Improvements

#9014 did make some decent improvements to read speed. When reading the dataset written above I get:

%timeit xr.backends.api.open_datatree(f"{bucket}/test.zarr", engine="zarr")
%timeit datatree.open_datatree(f"{bucket}/test.zarr", engine="zarr")

882 ms ± 47.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
3.47 s ± 86.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

We'll need similar optimizations on the write side. The fundamental issue is that DataTree.to_zarr relies on serial Dataset.to_zarr calls for each node:

xarray/xarray/core/datatree_io.py

Lines 153 to 171 in 12c690f

    
           for node in dt.subtree: 
        
               ds = node.to_dataset(inherited=False) 
        
               group_path = node.path 
        
               if ds is None: 
        
                   _create_empty_zarr_group(store, group_path, mode) 
        
               else: 
        
                   ds.to_zarr( 
        
                       store, 
        
                       group=group_path, 
        
                       mode=mode, 
        
                       encoding=encoding.get(node.path), 
        
                       consolidated=False, 
        
                       **kwargs, 
        
                   ) 
        
               if "w" in mode: 
        
                   mode = "a" 
        
           if consolidated: 
        
               consolidate_metadata(store)

This results in many fsspec calls to list dirs, check file existence, and put small metadata and attribute files in the bucket. Here's snakeviz on the example:

(The 8s block on the right is metadata consolidation)

Workaround

If your data is small enough to dump locally, this works great:

def to_zarr(dt, path):
    with TemporaryDirectory() as tmp_path:
        dt.to_zarr(tmp_path)
        fs.put(tmp_path, path, recursive=True)

Takes about 1s.

@jhamman

Dataset.to_zarr() is also very slow for Datasets with many variables.

We really need support for concurrency in writing Zarr metadata, which I believe may be coming in Zarr Python 3 (cc @jhamman)

@d-v-b

Yes, exactly right. We have build Zarr Python 3's interface around asyncIO. This will allow us to initialize hierarchies with much more concurrency than we do now. We haven't built the API for this yet into Zarr yet though but I think @d-v-b has experimented with this already.

@d-v-b

@d-v-b recently added an API in Zarr for concurrently creating multiple arrays/groups. This could now be applied to Xarray's Zarr backend:

@NathanCummings

I'd forgotten about this issue, but recently came across the same problem while looking at some of @NathanCummings DataTree data.

I made a simple benchmark for this problem, both to find out what needs fixing and to have a "before" for a "before-and-after" comparison.

(Code here: https://gist.github.com/TomNicholas/ed59c25a5c921ee8c971e211e26a7c83)

The benchmark writes increasingly large numbers of small variables to zarr, using LoggingStore to record all the zarr GET and SET operations, and using LatencyStore to artifically add latency to mimic the effect of writing to a remote cloud object store, and therefore measure the total time taken.

The results confirm what I suspected: there's a serial loop over both variables and over groups, which when you extrapolate the slope out to 100's or 1000's of groups and variables clearly leads to operations taking 10's of minutes even to write small volumes of data.

As @slevang noted above, this isn't surprising in that there is a serial loop over groups. But I am a bit surprised that Dataset.to_zarr also does a serial loop over variables. I also find it weird that even when doing more writes the number of GETs increases... There are clearly other opportunities for optimization here.

Ideally after optimization we would see the blue duration line be flat in all these cases. I think the green GETs line could also be made to be flat, but the red SETs line needs to go up linearly as we are writing more data.

I did also try with consolidated=True but it didn't make that much difference... I could try with an in-memory icechunk store but not sure if adding latency to every get/set makes sense then. Regardless this non-consolidated, non-icechunk high-latency store represents the worst case, so if we can make this fast all other writes to zarr should be fast too.

To actually fix this, I think we ideally want to refactor dump_to_store to accept a DataTree (or at least a dict[path, Dataset]), so that we have one top-level place that we can apply concurrency / parallelization across all groups and all variables.

To actually fix this, I think we ideally want to refactor dump_to_store to accept a DataTree (or at least a dict[path, Dataset]), so that we have one top-level place that we can apply concurrency / parallelization across all groups and all variables.

I am getting started on this refactor over in #10625

Navigation Menu

Repository navigation

`DataTree.to_zarr()` is very slow writing to high latency store #9455

What is your issue?

Test case

Potential Improvements

Workaround

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	for node in dt.subtree:
	ds = node.to_dataset(inherited=False)
	group_path = node.path
	if ds is None:
	_create_empty_zarr_group(store, group_path, mode)
	else:
	ds.to_zarr(
	store,
	group=group_path,
	mode=mode,
	encoding=encoding.get(node.path),
	consolidated=False,
	**kwargs,
	)
	if "w" in mode:
	mode = "a"

	if consolidated:
	consolidate_metadata(store)

DataTree.to_zarr() is very slow writing to high latency store #9455

Description

What is your issue?

Test case

Potential Improvements

Workaround

Activity

shoyer commented on Sep 8, 2024

jhamman commented on Sep 17, 2024

jhamman commented on Mar 20, 2025

TomNicholas commented on Aug 11, 2025

TomNicholas commented on Aug 11, 2025

shoyer commented on Aug 11, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

`DataTree.to_zarr()` is very slow writing to high latency store #9455