Set xarray "preferred_chunks" based on rasterio.block_shapes

[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]

Mind moving this to #253 ?

[scottyhq]

Mind moving this to https://github.com/corteva/rioxarray/issues/253 ?

🤦 sorry i overlooked that! Should I just copy/paste and close this one?

[snowman2]

sorry i overlooked that! Should I just copy/paste and close this one?

No worries. I overlook things all the time :smile:. Yes, please.