Add `metadata_only` param to `.to_zarr`?

[max-sixty]

Is your feature request related to a problem?

A leaf from #8245, which has a bullet:

compute=False is arguably a less-than-obvious kwarg meaning "write metadata". Maybe this should be a method, maybe it's a candidate for renaming? Or maybe make_template can be an abstraction over it

I've also noticed that for large arrays, running compute=False can take several minutes, despite the indexes being very small. I think this is because it's building a dask task graph — which is then discarded, since the array is written from different machines with the region pattern.

Describe the solution you'd like

Would introducing a metadata_only parameter to to_zarr help here:

• Better name
• No dask graph

Describe alternatives you've considered

No response

Additional context

No response

[shoyer]

Yes, is a great idea!

[jhamman]

+1, this is a really nice idea.

Related to this could also be a write-through cache of sorts. For high-latency stores (e.g. S3), synchronously populating the store metadata can really add up. If we knew we were only writing metadata, we could safely populate all the Zarr json objects then send them in one bulk write step.

The combination of these two features would be a lightning fast Zarr initialization routine 🚀

[max-sixty]

I came across #8343 recently — this seems to be a similar suggestion to what I was intending. Is that correct?

The challenge is that Xarray needs some way to represent the "schema" for the desired entire dataset. I'm very open to alternatives, but so far, the most convenient way to do this has been to load Dask arrays into an xarray.Dataset.

Is anyone more familiar with whether there is a cost to producing the dask task graph? I'm seeing .to_zarr(compute=False) take well over 1 minutes with large arrays with lots of chunks. And it's only writing very small metadata.

[shoyer]

Is anyone more familiar with whether there is a cost to producing the dask task graph? I'm seeing .to_zarr(compute=False) take well over 1 minutes with large arrays with lots of chunks. And it's only writing very small metadata.

If to_zarr(compute=False) is slow, that's more likely due to Xarray doing lots of sequential blocking IO operations. Once you've made the Dataset object the dask graphs have already been created.

[max-sixty]

If to_zarr(compute=False) is slow, that's more likely due to Xarray doing lots of sequential blocking IO operations. Once you've made the Dataset object the dask graphs have already been created.

This is an area I don't understand well, so obv I defer.

I had thought that it was driven by dask since that's most of the time is spent in dask/array/optimization.py. But maybe we're saying that this is because xarray is writing indexes during that time? (Though the indexes are fairly small, and it's spending lots of time in a fuse function?)

From py-spy:

  %Own   %Total  OwnTime  TotalTime  Function (filename)
  0.00%   0.00%   14.07s    16.42s   fuse (dask/optimization.py)
  0.00%   0.00%    8.15s     8.15s   reverse_dict (dask/core.py)
  0.00%   0.00%    3.17s     4.17s   make_blockwise_graph (dask/blockwise.py)
  0.00%   0.00%    2.81s     2.94s   _cull_dependencies (dask/blockwise.py)
 20.00%  25.00%    2.50s     3.04s   fuse_slice (dask/array/optimization.py)
  0.00%   0.00%    2.49s     2.49s   keys_in_tasks (dask/core.py)
  0.00%   0.00%    2.19s     2.90s   functions_of (dask/optimization.py)
  0.00%   0.00%    1.32s     9.20s   cull (dask/highlevelgraph.py)
  0.00%   0.00%   0.930s    48.16s   optimize (dask/array/optimization.py)

[dcherian]

Max I think you're right. In recent times, dask has a "lazy graph" (HighLevelGraph) that gets lowered down to an old-style graph expressed as dicts. That lowering is still slow and potentially whats happening here.

[dcherian]

Yeah here's the call to optimize in dask.array.store before it handles the compute kwarg

https://github.com/dask/dask/blob/1203b1bb6d52b1fb00d54656af4af8e6e35aa615/dask/array/core.py#L1169-L1171

[dcherian]

I really think we're just taking advantage of a side-effect by using compute=False for this use-case. I like the idea of a separate method, say initialize_zarr, for this.

[max-sixty]

Yes we could either do:

• .initialize_zarr / .to_zarr_metadata / xarray_beam.make_template
• or .to_zarr(metadata_only=True)

How difficult do we think this is? Is it something I can bang out in 45 mins or is it a bigger effort that requires more context?

[max-sixty]

FYI — for the moment, xr.ones_like(ds).to_zarr(compute=False) saves a few minutes each go! (though the data needs to be floats)

[shoyer]

FYI — for the moment, xr.ones_like(ds).to_zarr(compute=False) saves a few minutes each go! (though the data needs to be floats)

In xarray-beam we use zeros_like, which I believe works for any NumPy dtype.

[max-sixty]

How difficult do we think this is? Is it something I can bang out in 45 mins or is it a bigger effort that requires more context?

In xarray-beam we use zeros_like, which I believe works for any NumPy dtype.

Could this just be running xr.zero_like(ds).to_zarr(compute=False)?? Are there any data types that zarr supports which wouldn't be covered by zeros_like? Possibly strings...

[dcherian]

Possibly strings...

Seem to work!

ds = xr.Dataset(
    {
        "a": (("x", "y"), np.arange(20).reshape(4, 5)),
        "b": ("x", ["aa", "b", "c", "d"]),
        "c": ("x", np.array(["aa", "b", "c", "d"], dtype="S")),
        "time": ("t", pd.date_range("2001-01-01", periods=4, freq='D')),
    },
    coords={"x": np.arange(4), "y": [33, 44, 22, 11, 55]},
)