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)?

[jthielen]

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

Only indirectly, since this deals with duck Dask arrays (things like Pint that go between xarray and Dask) rather than Dask chunks, which CuPy would be. But, once this, #4212, hgrecco/pint#964, and dask/dask#6393 are all in place, then we can test if xarray( pint( dask( cupy ))) works automatically from it all or not.

[dcherian]

A couple of things came up in #4221

1. how do we ask a duck dask array to rechunk itself? pint seems to forward the .rechunk call but that isn't formalized anywhere AFAICT.
2. less important: should duck dask arrays cache their token somewhere? dask.array uses .name to do this and xarray uses that to check equality cheaply. We can use tokenize of course. But I'm wondering if it's worth asking duck dask arrays to cache their token as an optimization.

[mrocklin]

In Xarray we implemented the Dask collection spec. https://docs.dask.org/en/latest/custom-collections.html#the-dask-collection-interface

We might want to do that with Pint as well, if they're going to contain Dask things. That way Dask operations like dask.persist, dask.visualize, and dask.compute will work normally.

[mrocklin]

My guess is that we could steal the xarray.DataArray implementations over to Pint without causing harm.

[jthielen]

In Xarray we implemented the Dask collection spec. https://docs.dask.org/en/latest/custom-collections.html#the-dask-collection-interface

We might want to do that with Pint as well, if they're going to contain Dask things. That way Dask operations like dask.persist, dask.visualize, and dask.compute will work normally.

That's exactly what's been done in Pint (see hgrecco/pint#1129)! @dcherian's points go beyond just that and address what Pint hasn't covered yet through the standard collection interface.

[mrocklin]

That's exactly what's been done in Pint (see hgrecco/pint#1129)! @dcherian's points go beyond just that and address what Pint hasn't covered yet through the standard collection interface.

Ah, great. My bad.

how do we ask a duck dask array to rechunk itself? pint seems to forward the .rechunk call but that isn't formalized anywhere AFAICT.

I think that you would want to make a pint array rechunk method that called down to the dask array rechunk method. My guess is that this might come up in other situations as well.

less important: should duck dask arrays cache their token somewhere? dask.array uses .name to do this and xarray uses that to check equality cheaply. We can use tokenize of course. But I'm wondering if it's worth asking duck dask arrays to cache their token as an optimization.

I think that implementing the dask.base.normalize_token method should be fine. This will probably be very fast because you're probably just returning the name of the underlying dask array as well as the unit of the pint array/quatity. I don't think that caching would be necessary here.

It's also possible that we could look at the __dask_layers__ method to get this information. My memory is a bit fuzzy here though.

[dcherian]

Re:rechunk, this should be part of the spec I guess. We need this for DataArray.chunk().

xarray does do some automatic rechunking in variable.py. But this comment:

            # chunked data should come out with the same chunks; this makes
            # it feasible to combine shifted and unshifted data
            # TODO: remove this once dask.array automatically aligns chunks

suggest that we could delete that automatic rechunking today.

This will probably be very fast because you're probably just returning the name of the underlying dask array as well as the unit of the pint array/quatity.

ah yes, we can rely on the underlying array library to optimize this.

[mrocklin]

Dask collections tokenize quickly. We just use the name I think.