xr.map_blocks vs dask.array.map_blocks wrapping numpy function

[aaronspring]

issue arised in zarr-developers/numcodecs#299

I want to apply a numpy function

from numcodecs.bitround import BitRound
def _np_bitround(data, keepbits):
    """Bitround for Arrays."""
    codec = BitRound(keepbits=keepbits)
    data = data.copy()  # otherwise overwrites the input
    encoded = codec.encode(data)
    return codec.decode(data)

to xr objects with dask.

My observations:

• Works with xr.apply_ufunc(dask="parallelized"):
  xr.apply_ufunc(_np_bitround, ds, 7, dask='parallelized').compute() 👍
• Fails with xr.apply_ufunc(dask="allowed") at the compute() step:

py:57, in BitRound.encode(self, buf)
     51 def encode(self, buf):
     52     """Create int array by rounding floating-point data
     53 
     54     The itemsize will be preserved, but the output should be much more
     55     compressible.
     56     """
---> 57     a = ensure_ndarray_like(buf)
     58     if not a.dtype.kind == "f" or a.dtype.itemsize > 8:
     59         raise TypeError("Only float arrays (16-64bit) can be bit-rounded")

File /work/mh0727/m300524/conda-envs/bitinfo/lib/python3.10/site-packages/numcodecs/compat.py:42, in ensure_ndarray_like(buf)
     38     raise TypeError("array.array with char or unicode type is not supported")
     39 else:
     40     # N.B., first take a memoryview to make sure that we subsequently create a
     41     # numpy array from a memory buffer with no copy
---> 42     mem = memoryview(buf)
     43     # instantiate array from memoryview, ensures no copy
     44     buf = np.array(mem, copy=False)

TypeError: memoryview: a bytes-like object is required, not 'Array'

• Fails with xr.map_blocks(dask="allowed") at the compute() step:

File /work/mh0727/m300524/conda-envs/bitinfo/lib/python3.10/site-packages/numcodecs/bitround.py:57, in BitRound.encode(self, buf)
     51 def encode(self, buf):
     52     """Create int array by rounding floating-point data
     53 
     54     The itemsize will be preserved, but the output should be much more
     55     compressible.
     56     """
---> 57     a = ensure_ndarray_like(buf)
     58     if not a.dtype.kind == "f" or a.dtype.itemsize > 8:
     59         raise TypeError("Only float arrays (16-64bit) can be bit-rounded")

File /work/mh0727/m300524/conda-envs/bitinfo/lib/python3.10/site-packages/numcodecs/compat.py:42, in ensure_ndarray_like(buf)
     38     raise TypeError("array.array with char or unicode type is not supported")
     39 else:
     40     # N.B., first take a memoryview to make sure that we subsequently create a
     41     # numpy array from a memory buffer with no copy
---> 42     mem = memoryview(buf)
     43     # instantiate array from memoryview, ensures no copy
     44     buf = np.array(mem, copy=False)

TypeError: memoryview: a bytes-like object is required, not 'Dataset'

but works for dask.array.map_blocks: 👍
import dask.array as da; da.map_blocks(_np_bitround, ds[v].data, keepbits=7, dtype="float32").compute()

So there might be something in xr.map_blocks not being properly transferred to dask.array.map_blocks or I miss an important keyword.

Additional context: https://stackoverflow.com/questions/51736172/whats-the-difference-between-dask-parallelized-and-dask-allowed-in-xarrays-app

[dcherian]

Please see https://docs.xarray.dev/en/stable/user-guide/dask.html#automatic-parallelization-with-apply-ufunc-and-map-blocks

xr.map_blocks does not wrap dask.array.map_blocks. xr.apply_ufunc(..., dask="parallelized") wraps dask.array.apply_gufunc which ends up being equivalent to dask.array.map_blocks.

The underlying function you're using is designed to work with numpy arrays. Notice np.array(mem)

[aaronspring]

ah ok. now I understand why dask="allowed" and xr.map_blocks fail. And dask.array.map_blocks doesnt fail and ends up being something like xr.apply_ufunc(..., dask="parallelized").

[dcherian]

Yes exactly. Do you have suggestions for how to improve the documentation on those point?

[aaronspring]

RTFM in big letters ;) I read the docs before but never understood these differences before having this first user case for myself. https://docs.xarray.dev/en/stable/user-guide/dask.html#automatic-parallelization-with-apply-ufunc-and-map-blocks is better clear

[aaronspring]

"‘allowed’: pass dask arrays directly on to func. Prefer this option if func natively supports dask arrays." Maybe add that numpy functions likely fail. But I guess not all, but if they contain np.array() stuff.