Default chunking when reading and writing with dask

[snowman2]

trafficstars

Related to: https://github.com/pangeo-data/cog-best-practices/issues/2

https://github.com/pangeo-data/cog-best-practices/blob/main/4-threads-vs-async.ipynb

Sounds like it would be worth digging into and seeing what could be improved.

[snowman2]

For reference, current location of chunking code:

https://github.com/corteva/rioxarray/blob/ce1194acb4ababbb4c6570d72cf64aa2760c0fa2/rioxarray/_io.py#L650-L688

[scottyhq]

It could be nice to take advantage of known raster tiling (rasterio.block_shapes) to assign dask chunks when opening tiled datasets. It seems straightforward to implement with xarray's backend preferred chunk sizes.

Setting this could also avoid common errors with exploding task graphs when users rely on defaults (chunks=True) or reading way more data than necessary (chunks={}).

According to xarray docs for open_dataset/open_dataarray, "chunks={} loads the dataset with dask using engine preferred chunks if exposed by the backend, otherwise with a single chunk for all arrays. chunks='auto' will use dask auto chunking". Below is an example with some public data (18287x18460 with 256x256 tiling):

import xarray as xr
import rasterio
url = 'https://capella-open-data.s3.amazonaws.com/data/2022/4/21/CAPELLA_C06_SP_GEO_HH_20220421064500_20220421064519/CAPELLA_C06_SP_GEO_HH_20220421064500_20220421064519_preview.tif'

with rasterio.open(url) as src:
    print(src.block_shapes) # [(256, 256)]

# This is okay, but not optimal due to misalignment with 256 block_shapes
da = xr.open_dataarray(url, chunks='auto', engine='rasterio')
# da.chunksizes 
# Frozen({'band': (1,), 'y': (9087, 9087, 113), 'x': (3692, 3692, 3692, 3692, 3692)})

# Warning! This is the same as chunks=1, which explodes the number of tasks
da = xr.open_dataarray(url, chunks=True, engine='rasterio')

# Warning! You'll be reading all pixel values even if you just want 1
xr.open_dataarray(url, chunks={}, engine='rasterio')

There are many considerations for optimal chunks (see dask blog), but a good default size could be determined by ~100Mb and a multiple of 256x256?

[snowman2]

@scottyhq. that sounds like a great place to start for improving the default chunking. Thanks for sharing your thoughts on this :+1: