Support for duck Dask Arrays

[jthielen]

#525 (comment) raised the idea of adding "duck Dask Array" support to xarray as a way to handle xarray > Pint Quantity > Dask Array wrapping in a way that still allowed for most of xarray's Dask integration to work properly. With @rpmanser working on implementing the Dask collection interface in Pint (hgrecco/pint#1129), I thought it best to elevate this to its own issue to track progress and discuss implementation on xarray's side (since hopefully @rpmanser or I can get started on it soon).

Two initial (and intertwined) discussion points that I'd like to bring up (xref hgrecco/pint#1129 (comment)):

• How should xarray check for a duck Dask Array?
• Is it acceptable for a Pint Quantity to always have the Dask collection interface defined (i.e., be a duck Dask array), even when its magnitude (what it wraps) is not a Dask Array?

cc @keewis, @shoyer, @crusaderky

[shoyer]

Maybe something like this would work?

def is_duck_dask_array(x):
  return getattr(x, 'chunks', None) is not None

xarray.DataArray would pass this test (chunks is either None for non-dask arrays or a tuple for dask arrays), so this would be consistent with what we already do.

[jthielen]

Maybe something like this would work?

def is_duck_dask_array(x):
  return getattr(x, 'chunks', None) is not None

xarray.DataArray would pass this test (chunks is either None for non-dask arrays or a tuple for dask arrays), so this would be consistent with what we already do.

That would be a straightforward solution to both problems! A Pint Quantity containing a Dask Array passes along the chunks attribute from the Dask Array, and a Pint Quantity containing something else will raise an AttributeError. Unless there are other objections, I'll see what it will take to swap out the existing Dask checks for this in the xarray internals and hopefully get around to a PR (after I get some MetPy stuff done first).

[shoyer]

It might also make sense to check for one or more of the special dask collection attributes (__dask_graph__, __dask_keys__, etc)

[crusaderky]

Is it acceptable for a Pint Quantity to always have the Dask collection interface defined (i.e., be a duck Dask array), even when its magnitude (what it wraps) is not a Dask Array?

I think there are already enough headaches with __iter__ being always defined and confusing libraries such as pandas (hgrecco/pint#1128).
I don't see why pint should be explicitly aware of dask (except in unit tests)? It should only deal with generic NEP18-compatible libraries (numpy, dask, sparse, cupy, etc.).

How should xarray check for a duck Dask Array?

We should ask the dask team to formalize what defines a "dask-array-like", like they already did with dask collections, and implement their definition in xarray.
I'd personally make it "whatever defines a numpy-array-like AND has a chunks method AND the chunks method returns a tuple".

[jthielen]

I think there are already enough headaches with __iter__ being always defined and confusing libraries such as pandas (hgrecco/pint#1128).
I don't see why pint should be explicitly aware of dask (except in unit tests)? It should only deal with generic NEP18-compatible libraries (numpy, dask, sparse, cupy, etc.).

Since Pint wraps Dask, in order to leverage Dask Array functionality on Pint Quantities, we need to have the Dask collection interface available. In a sense, Pint needs special handling for Dask like xarray Variables do since they both can be upcast types of Dask Array. Implicitly passing through attributes (how Pint handles special methods/attributes of downcast types in general) from the wrapped Dask Array is not sufficient, however, because the finalizers have to rewrap with Quantity (see https://github.com/hgrecco/pint/pull/1129/files#diff-d9924213798d0fc092b8cff13928d747R1947-R1950), hence the explicit awareness of Dask being needed in Pint.

We should ask the dask team to formalize what defines a "dask-array-like", like they already did with dask collections, and implement their definition in xarray.

Done! See dask/dask#6385.

[jthielen]

Based on @mrocklin's comment in dask/dask#6385, the plan will be to check for duck Dask Arrays with dask.base.is_dask_collection along with xarray's previously used duck array check. This works properly with Pint Quantities as implemented in hgrecco/pint#1129 (returning True if the Pint Quantity contains a Dask Array, and False if it does not).

[rpmanser]

I can go ahead with putting together a PR for this. Before I do so, I'd like to clarify what is expected.

• Implement the is_duck_dask_array() function
• In that implementation, use dask.base.is_dask_collection() and the existing duck array check(s)
• Replace isinstance(x, dask.array.Array) checks with the new is_dask_duck_array() function

I searched for existing duck array checks in xarray and nothing immediately obvious to me is showing up. It looks like a check for __array_function__ is inappropriate based on discussion in #3917. Could someone point out the proper duck array check?

[dcherian]

We have https://github.com/pydata/xarray/blob/master/xarray/core/pycompat.py which defines dask_array_type and sparse_array_type and then use isinstance(da, dask_array_type) in a bunch of places (e.g. duck_array_ops).

re duck array check: @keewis added this recently

xarray/xarray/core/utils.py

Lines 250 to 253 in f3ca63a

	def is_array_like(value: Any) -> bool:
	return (
	hasattr(value, "ndim") and hasattr(value, "shape") and hasattr(value, "dtype")
	)

[jthielen]

@rpmanser That sounds like a good plan to me at least, but it would be great if any of the xarray maintainers would be able to chime in. Also, thanks again for being willing to work on this while I try working on dask/dask#4583. The hidden 4th step is of course testing--primarily that this doesn't break existing functionality, but also that it works for duck Dask Arrays other than Dask Arrays themselves (not sure if Pint Quantities in upcoming v0.15 or a mocked class would be better).

Also, thank you @dcherian for pointing out those checks, you found them faster than I did!

[dopplershift]

Does/should any of this also consider #4212 (CuPy)?