Incrementally write to Zarr

[davidbrochart]

A Dataset's to_zarr with compute=False allows for writing the array lazily. I'd like to use that to incrementally write data to the Zarr store, e.g.:

import numpy as np
import xarray as xr

x = np.linspace(0, 10, 100)
y = np.linspace(50, 40, 100)
data = np.random.randn(100, 100)
da = xr.DataArray(data, coords={"y": y, "x": x}, dims=["y", "x"])
da = da.chunk((10, 10))

ds = xr.Dataset({"da": da})

delayed = ds.to_zarr("ds", compute=False)

da[:10, :10] = 1

delayed.compute()

But this writes all the chunks, while I'd like to write only one.
Is there a better way?

[spencerkclark]

Have you looked into the region keyword argument of to_zarr? I think it should allow you to do what you are looking for; see this section of the documentation for an example.

[davidbrochart]

Thanks @spencerkclark, indeed region seems to be what I need. But I'm confused as how to use it.
With:

ds.to_zarr("ds", region={"x": slice(0, 10), "y": slice(0, 10)})

I get:

GroupNotFoundError: group not found at path ''

/tmp/ipykernel_87506/3355768128.py:1: RuntimeWarning: Failed to open Zarr store with consolidated metadata, falling back to try reading non-consolidated metadata. This is typically much slower for opening a dataset. To silence this warning, consider:
1. Consolidating metadata in this existing store with zarr.consolidate_metadata().
2. Explicitly setting consolidated=False, to avoid trying to read consolidate metadata, or
3. Explicitly setting consolidated=True, to raise an error in this case instead of falling back to try reading non-consolidated metadata.
  ds.to_zarr("ds", region={"x": slice(0, 10), "y": slice(0, 10)})
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/backends/zarr.py:384, in ZarrStore.open_group(cls, store, mode, synchronizer, group, consolidated, consolidate_on_close, chunk_store, storage_options, append_dim, write_region, safe_chunks, stacklevel)
    383 try:
--> 384     zarr_group = zarr.open_consolidated(store, **open_kwargs)
    385 except KeyError:

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/convenience.py:1304, in open_consolidated(store, metadata_key, mode, **kwargs)
   1303 # setup metadata store
-> 1304 meta_store = ConsolidatedStoreClass(store, metadata_key=metadata_key)
   1306 # pass through

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/storage.py:2853, in ConsolidatedMetadataStore.__init__(self, store, metadata_key)
   2852 # retrieve consolidated metadata
-> 2853 meta = json_loads(self.store[metadata_key])
   2855 # check format of consolidated metadata

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/storage.py:1067, in DirectoryStore.__getitem__(self, key)
   1066 else:
-> 1067     raise KeyError(key)

KeyError: '.zmetadata'

During handling of the above exception, another exception occurred:

GroupNotFoundError                        Traceback (most recent call last)
Input In [5], in <cell line: 1>()
----> 1 ds.to_zarr("ds", region={"x": slice(0, 10), "y": slice(0, 10)})

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/core/dataset.py:2060, in Dataset.to_zarr(self, store, chunk_store, mode, synchronizer, group, encoding, compute, consolidated, append_dim, region, safe_chunks, storage_options)
   1950 """Write dataset contents to a zarr group.
   1951 
   1952 Zarr chunks are determined in the following way:
   (...)
   2056     The I/O user guide, with more details and examples.
   2057 """
   2058 from ..backends.api import to_zarr
-> 2060 return to_zarr(  # type: ignore
   2061     self,
   2062     store=store,
   2063     chunk_store=chunk_store,
   2064     storage_options=storage_options,
   2065     mode=mode,
   2066     synchronizer=synchronizer,
   2067     group=group,
   2068     encoding=encoding,
   2069     compute=compute,
   2070     consolidated=consolidated,
   2071     append_dim=append_dim,
   2072     region=region,
   2073     safe_chunks=safe_chunks,
   2074 )

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/backends/api.py:1597, in to_zarr(dataset, store, chunk_store, mode, synchronizer, group, encoding, compute, consolidated, append_dim, region, safe_chunks, storage_options)
   1595     already_consolidated = False
   1596     consolidate_on_close = consolidated or consolidated is None
-> 1597 zstore = backends.ZarrStore.open_group(
   1598     store=mapper,
   1599     mode=mode,
   1600     synchronizer=synchronizer,
   1601     group=group,
   1602     consolidated=already_consolidated,
   1603     consolidate_on_close=consolidate_on_close,
   1604     chunk_store=chunk_mapper,
   1605     append_dim=append_dim,
   1606     write_region=region,
   1607     safe_chunks=safe_chunks,
   1608     stacklevel=4,  # for Dataset.to_zarr()
   1609 )
   1611 if mode in ["a", "r+"]:
   1612     _validate_datatypes_for_zarr_append(zstore, dataset)

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/backends/zarr.py:401, in ZarrStore.open_group(cls, store, mode, synchronizer, group, consolidated, consolidate_on_close, chunk_store, storage_options, append_dim, write_region, safe_chunks, stacklevel)
    385     except KeyError:
    386         warnings.warn(
    387             "Failed to open Zarr store with consolidated metadata, "
    388             "falling back to try reading non-consolidated metadata. "
   (...)
    399             stacklevel=stacklevel,
    400         )
--> 401         zarr_group = zarr.open_group(store, **open_kwargs)
    402 elif consolidated:
    403     # TODO: an option to pass the metadata_key keyword
    404     zarr_group = zarr.open_consolidated(store, **open_kwargs)

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/hierarchy.py:1347, in open_group(store, mode, cache_attrs, synchronizer, path, chunk_store, storage_options, zarr_version)
   1345         if contains_array(store, path=path):
   1346             raise ContainsArrayError(path)
-> 1347         raise GroupNotFoundError(path)
   1349 elif mode == 'w':
   1350     init_group(store, overwrite=True, path=path, chunk_store=chunk_store)

GroupNotFoundError: group not found at path ''

This is probably because by default mode="r+" when region is passed. If I pass mode="a", I get:

ValueError: parameter 'value': expected array with shape (10,), got (100,)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [5], in <cell line: 1>()
----> 1 ds.to_zarr("ds", region={"x": slice(0, 10), "y": slice(0, 10)}, mode="a")

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/core/dataset.py:2060, in Dataset.to_zarr(self, store, chunk_store, mode, synchronizer, group, encoding, compute, consolidated, append_dim, region, safe_chunks, storage_options)
   1950 """Write dataset contents to a zarr group.
   1951 
   1952 Zarr chunks are determined in the following way:
   (...)
   2056     The I/O user guide, with more details and examples.
   2057 """
   2058 from ..backends.api import to_zarr
-> 2060 return to_zarr(  # type: ignore
   2061     self,
   2062     store=store,
   2063     chunk_store=chunk_store,
   2064     storage_options=storage_options,
   2065     mode=mode,
   2066     synchronizer=synchronizer,
   2067     group=group,
   2068     encoding=encoding,
   2069     compute=compute,
   2070     consolidated=consolidated,
   2071     append_dim=append_dim,
   2072     region=region,
   2073     safe_chunks=safe_chunks,
   2074 )

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/backends/api.py:1637, in to_zarr(dataset, store, chunk_store, mode, synchronizer, group, encoding, compute, consolidated, append_dim, region, safe_chunks, storage_options)
   1635 writer = ArrayWriter()
   1636 # TODO: figure out how to properly handle unlimited_dims
-> 1637 dump_to_store(dataset, zstore, writer, encoding=encoding)
   1638 writes = writer.sync(compute=compute)
   1640 if compute:

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/backends/api.py:1257, in dump_to_store(dataset, store, writer, encoder, encoding, unlimited_dims)
   1254 if encoder:
   1255     variables, attrs = encoder(variables, attrs)
-> 1257 store.store(variables, attrs, check_encoding, writer, unlimited_dims=unlimited_dims)

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/backends/zarr.py:571, in ZarrStore.store(self, variables, attributes, check_encoding_set, writer, unlimited_dims)
    568     self.set_attributes(attributes)
    569     self.set_dimensions(variables_encoded, unlimited_dims=unlimited_dims)
--> 571 self.set_variables(
    572     variables_encoded, check_encoding_set, writer, unlimited_dims=unlimited_dims
    573 )
    574 if self._consolidate_on_close:
    575     zarr.consolidate_metadata(self.zarr_group.store)

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/backends/zarr.py:650, in ZarrStore.set_variables(self, variables, check_encoding_set, writer, unlimited_dims)
    647     zarr_array.resize(new_shape)
    649 region = tuple(write_region[dim] for dim in dims)
--> 650 writer.add(v.data, zarr_array, region)

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/xarray/backends/common.py:156, in ArrayWriter.add(self, source, target, region)
    154 else:
    155     if region:
--> 156         target[region] = source
    157     else:
    158         target[...] = source

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/core.py:1353, in Array.__setitem__(self, selection, value)
   1351     self.vindex[selection] = value
   1352 else:
-> 1353     self.set_basic_selection(pure_selection, value, fields=fields)

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/core.py:1448, in Array.set_basic_selection(self, selection, value, fields)
   1446     return self._set_basic_selection_zd(selection, value, fields=fields)
   1447 else:
-> 1448     return self._set_basic_selection_nd(selection, value, fields=fields)

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/core.py:1748, in Array._set_basic_selection_nd(self, selection, value, fields)
   1742 def _set_basic_selection_nd(self, selection, value, fields=None):
   1743     # implementation of __setitem__ for array with at least one dimension
   1744 
   1745     # setup indexer
   1746     indexer = BasicIndexer(selection, self)
-> 1748     self._set_selection(indexer, value, fields=fields)

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/core.py:1776, in Array._set_selection(self, indexer, value, fields)
   1774     if not hasattr(value, 'shape'):
   1775         value = np.asanyarray(value)
-> 1776     check_array_shape('value', value, sel_shape)
   1778 # iterate over chunks in range
   1779 if not hasattr(self.store, "setitems") or self._synchronizer is not None \
   1780    or any(map(lambda x: x == 0, self.shape)):
   1781     # iterative approach

File ~/mambaforge/envs/xarray/lib/python3.10/site-packages/zarr/util.py:547, in check_array_shape(param, array, shape)
    544     raise TypeError('parameter {!r}: expected an array-like object, got {!r}'
    545                     .format(param, type(array)))
    546 if array.shape != shape:
--> 547     raise ValueError('parameter {!r}: expected array with shape {!r}, got {!r}'
    548                      .format(param, shape, array.shape))

ValueError: parameter 'value': expected array with shape (10,), got (100,)

Indeed my array has shape (100, 100), I thought the region would also apply to ds. If not, it means I also have to slice my dataset, which will write a region of my coordinates to Zarr, it doesn't really makes sense.
Am I missing something?

[dcherian]

Yes you have to slice your dataset (potentially these datasets are being written from separate processes or separate workers): https://docs.xarray.dev/en/stable/user-guide/io.html#appending-to-existing-zarr-stores

[davidbrochart]

Thanks @dcherian, I should have read the docs!

[davidbrochart]

I noticed the initial .to_zarr() on the dataset created with the Dask array takes more than 2 minutes to process, although it's supposed to only write metadata.
My array has shape (524288, 524288) and chunks (256, 256), but it shouldn't even matter, right? Or maybe I'm running into some Dask overhead (4 million chunks)?

[dcherian]

maybe #6084