Extreamly slow write of rechunked xr.Dataset with xr.Dataset.to_zarr()

[mliukis]

Hello,

I am reading Zarr store from S3 bucket into xr.Dataset. I rechunk dataset with 2 dimensions chunk sizes set to 1:

dask_chunks = {'mid_date': 11117, 'x': 1, 'y': 1}

and want to save it to local file. It seems that to_zarr() is extremely slow when chunk size is small. Did anybody run into a similar issue?
It takes extremely long time (4.5hours so far to write 5.8Mb out of 5.4Gb) to write such chunked xr.Dataset to Zarr.

It is faster if I set dask_chunks = {'mid_date': 11117, 'x': 10, 'y': 10}: wrote 980Mb out of 5.4Gb in 1.5 hours. Still it took more than 8 hours to write the whole dataset to disk.

It has reasonable runtime (tens of minutes) if I change chunk size to a larger value:

dask_chunks = {'mid_date': 11117, 'x': 250, 'y': 250}

Just wonder if it's xarray, dask or zarr issue.

Also does anybody know if selected size of chunking affects access speed of the data? We are only interested in time series of one spacial point (x, y) of our dataset at a time. Would accessing such single spacial point time series be affected by the chunk size of our dataset it's stored in on the file system?

import s3fs 
import xarray as xr
import zarr

# Encoding settings for Zarr format
COMPRESSION_ZARR = zarr.Blosc(cname='zlib', clevel=2, shuffle=1)

ENCODING_ZARR = {
    'interp_mask':               {'_FillValue': 0.0, 'dtype': 'ubyte'},
    'flag_stable_shift':         {'_FillValue': None, 'dtype': 'long'},
    'chip_size_height':          {'_FillValue': 0.0, 'dtype': 'ushort'},
    'chip_size_width':           {'_FillValue': 0.0, 'dtype': 'ushort'},
    'stable_count_slow':         {'_FillValue': None, 'dtype': 'long'},
    'stable_count_mask':         {'_FillValue': None, 'dtype': 'long'},
    'v_error':                   {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'v':                         {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'vx':                        {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'vx_error':                  {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vx_error_mask':             {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vx_error_modeled':          {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vx_error_slow':             {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vx_stable_shift':           {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vx_stable_shift_slow':      {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vx_stable_shift_mask':      {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vy':                        {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'vy_error':                  {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vy_error_mask':             {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vy_error_modeled':          {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vy_error_slow':             {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vy_stable_shift':           {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vy_stable_shift_slow':      {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vy_stable_shift_mask':      {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'va':                        {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'va_error':                  {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'va_error_mask':             {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'va_error_modeled':          {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'va_error_slow':             {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'va_stable_shift':           {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'va_stable_shift_slow':      {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'va_stable_shift_mask':      {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vr':                        {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'vr_error':                  {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vr_error_mask':             {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vr_error_modeled':          {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vr_error_slow':             {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vr_stable_shift':           {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vr_stable_shift_slow':      {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vr_stable_shift_mask':      {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vxp':                       {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'vxp_error':                 {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vxp_error_mask':            {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vxp_error_modeled':         {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vxp_error_slow':            {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vxp_stable_shift':          {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vxp_stable_shift_slow':     {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vxp_stable_shift_mask':     {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vyp':                       {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'vyp_error':                 {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vyp_error_mask':            {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vyp_error_modeled':         {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vyp_error_slow':            {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vyp_stable_shift':          {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vyp_stable_shift_slow':     {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vyp_stable_shift_mask':     {'_FillValue': -32767.0, 'dtype': 'double', 'compressor': COMPRESSION_ZARR},
    'vp':                        {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'vp_error':                  {'_FillValue': -32767.0, 'dtype': 'short', 'compressor': COMPRESSION_ZARR},
    'acquisition_date_img1':     {'_FillValue': None, 'units': 'days since 1970-01-01'},
    'acquisition_date_img2':     {'_FillValue': None, 'units': 'days since 1970-01-01'},
    'date_center':               {'_FillValue': None, 'units': 'days since 1970-01-01'},
    'mid_date':                  {'_FillValue': None, 'units': 'days since 1970-01-01'},
    'roi_valid_percentage':      {'_FillValue': None},
    'satellite_img1':            {'_FillValue': None},
    'satellite_img2':            {'_FillValue': None},
    'mission_img1':              {'_FillValue': None},
    'mission_img2':              {'_FillValue': None},
    'sensor_img1':               {'_FillValue': None},
    'sensor_img2':               {'_FillValue': None},
    'autoRIFT_software_version': {'_FillValue': None},
    'date_dt':                   {'_FillValue': None}
}

cube_dir = "s3://its-live-data.jpl.nasa.gov/datacubes/v1/N60W040/ITS_LIVE_vel_EPSG3413_G0120_X-150000_Y-2250000.zarr"
fixed_file = "ITS_LIVE_vel_EPSG3413_G0120_X-150000_Y-2250000_rechunked.zarr"

s3_in = s3fs.S3FileSystem(anon=True)
cube_store = s3fs.S3Map(root=cube_dir, s3=s3_in, check=False)

dask_chunks = {'mid_date': 11117, 'x': 1, 'y': 1}

with xr.open_dataset(cube_store, decode_timedelta=False, engine='zarr', consolidated=True) as ds:
    # Re-chunk data to max length for mid_date
    chunked_ds = ds.chunk(dask_chunks)
    chunked_ds.to_zarr(fixed_file, encoding=ENCODING_ZARR, consolidated=True)

xr.show_versions()

INSTALLED VERSIONS

commit: None
python: 3.8.8 | packaged by conda-forge | (default, Feb 20 2021, 16:12:38)
[Clang 11.0.1 ]
python-bits: 64
OS: Darwin
OS-release: 19.6.0
machine: x86_64
processor: i386
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.10.6
libnetcdf: 4.7.4

xarray: 0.19.0
pandas: 1.1.4
numpy: 1.20.2
scipy: None
netCDF4: 1.5.4
pydap: None
h5netcdf: 0.8.1
h5py: 3.1.0
Nio: None
zarr: 2.6.1
cftime: 1.4.1
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: 2.30.0
distributed: 2.30.1
matplotlib: 3.3.2
cartopy: None
seaborn: None
numbagg: None
pint: None
setuptools: 52.0.0.post20210125
pip: 20.2.4
conda: None
pytest: None
IPython: 7.22.0
sphinx: None