Correct initialization of an "empty" Zarr dataset for region-based writing.

[mdespriee]

We're dealing with quite a big dataset that doesn't fit into memory. We're receiving it in small pieces, and we aim to incrementally create a Zarr store by writing each part to a specific region.
We've designed this general flow:

1- Create an "empty" Dataset with the correct dimensions.

empty = xr.Dataset(
   data_vars={'foo': (['x','y','t'], np.full((20,20,20), np.nan)) },
   coords={'x':np.arange(20), 'y':np.arange(20), 't':np.arange(20)}
)

2- Store it in Zarr.

empty.to_zarr('my_dataset.zarr')

3- Then, for each part, write using the region.

part = xr.Dataset(
  {'foo': (['x','y','t'], np.full((5,5,5), 10)) },
   coords={'x':np.arange(5,10), 'y':np.arange(5,10), 't':np.arange(5,10)},
)
part.to_zarr('my_dataset.zarr', region={"x": "auto", "y": "auto", "t": "auto"})

However, we're facing an issue with the very first part. The np.full(...) command creates a full-sized array that doesn't fit in memory. We've tried other creation methods, but none are working with region-based writing.

We're likely missing something. What's the appropriate approach to prepare our Zarr store without materializing the full array ?

[dcherian]

See https://docs.xarray.dev/en/stable/user-guide/io.html#modifying-existing-zarr-stores

Something like xr.zeros_like(ds).to_zarr(..., compute=False) will initialize the store without allocating a bunch of memory.

[mdespriee]

Thanks. But shouldn't ds be initialized first ?

[dcherian]

If you don't have ds then you can create a lazy "schema" with dask. This is shown in the doc link I put in above.

[mdespriee]

got it, the dask part what was missing. thanks

[mdespriee]

Ah, I found out that dask.array.full can be used instead of np.full to initialize lazily.
And I forgot compute=False

In summary

empty = xr.Dataset(
   data_vars={'foo': (['x','y','t'], dask.array.full((20,20,20), np.nan)) },
   coords={'x':np.arange(20), 'y':np.arange(20), 't':np.arange(20)}
)
empty.to_zarr('my_dataset.zarr', compute=False)
part = xr.Dataset(
  {'foo': (['x','y','t'], np.full((5,5,5), 10)) },
   coords={'x':np.arange(5,10), 'y':np.arange(5,10), 't':np.arange(5,10)},
)
part.to_zarr('my_dataset.zarr', region={"x": "auto", "y": "auto", "t": "auto"})

[dcherian]

Great! We'd happily take any suggestions for improving the docs for this use-case.