Hooks for XArray operations

[hameerabbasi]

In hope of cleaner dask and sparse support (pydata/sparse#1), I wanted to suggest hooks for XArray operations.

Something like the following:

try:
    import dask.array as da
    xarray.hooks.register('nansum', da.array, da.nansum)
    ...
except ImportError:
    pass


try:
    import sparse.SparseArray
    xarray.hooks.register('nansum', sparse.SparseArray, sparse.nansum)
    ...
except ImportError:
    pass

Functions would work something like the following:
(the register would fall back to Numpy if nothing is found)

• Check type of first/primary argument.
• Check register for function.
• Call function

I would argue that this should be in Numpy, but it's a huge project to put it there.

[fujiisoup]

Thanks for leading the development of sparse.
I'm looking forward to see it in xarray:)

Currently, our logic to support dask.array and numpy.ndarray is hard-coded everywhere.
For example, we have many computation paths for nansum,
dask.Array, np.ndarray with bottleneck, bare np.ndarray and we use our in-house implementation for object-type arrays.
The easiest way to support sparse might be to add a specific path for sparse by hard-coding again,
but it is less flexible.

Do we need to be capable of supporting other objects for future extension?
If so, we may need to start from (heavy) refactoring.

@shoyer,
Could you give any suggestion?
I am personally interested in helping this, but I may need to decide the direction first.

[hameerabbasi]

Then I would suggest something like the following for hooks (omitting imports):

# Registered in order of priority
xarray.interfaces.register('DaskArray', lambda ar: isinstance(ar, da.array))
xarray.hooks.register('nansum', 'DaskArray', da.nansum)

xarray.interfaces.register('SparseArray', lambda ar: isinstance(ar, sparse.SparseArray))
xarray.hooks.register('nansum', 'SparseArray', sparse.nansum)

And then, in code, call the appropriate nansum instead of np.nansum:

nansum = xarray.hooks.get(arr, 'nansum')

If you need help, I'd be willing to give it. :-) But I'm not a user of XArray, so I don't really understand the use-cases or codebase.

[shoyer]

Do we need to be capable of supporting other objects for future extension?
If so, we may need to start from (heavy) refactoring.

For two array backends, it didn't make sense to write an abstraction layer for this, in part because it wasn't clear what we needed. But for three examples, it probably does -- that's the point where shared use cases become clear. Undoubtedly, there will be other cases in the future where users will want to extend xarray to handle new array types (arrays with units come to mind).

For implementing these overloads/functions, there are various possible solutions. Our current ad-hoc system is similar to what @hameerabbasi suggests -- we check the type of the first argument and use that to dispatch to an appropriate function. This has the advantage of being easy to implement for a known set of types, but a single dispatch order is not very extensible -- it's impossible to anticipate every third-party class. Recently, NumPy has moved away from this (e.g., with __array_ufunc__).

One appealing option is to make use of @mrocklin's multipledispatch library, which was originally developed for Blaze and is still in active use. Possible concerns:

1. Performance. Import times need to be fast, and I know this is something that multipledispatch can sometimes struggle with. My guess is that this wouldn't be a problem for us, since we can rely on other dispatch mechanisms most operations (including __array_ufunc__ and Python's builtin arithmetic overrides).
2. Dispatch for stack/concatenate: How do we handle dispatching for functions that take a list of arrays? e.g., if a list of arrays has contains any dask arrays, we need to use dask. Ideally, we would resolve the type of an object like [np.array(...), np.array(...), ..., da.Array(...)] to a mixed type like List[Union[np.ndarray, da.Array]], for which an override could be implemented.
3. Dispatch for the first argument(s) only: This is a minor point, but some functions don't need to be dispatched on all of their arguments, e.g., sum() only really needs to dispatch on the array types but can pass other arguments like axis directly on. I suppose could simply annotate extra position arguments with object, but this will get annoying for multiple optional arguments which would all need separate implementations (if I understand multipledispatch correctly).

[mrocklin]

Import times on multipledispatch have improved thanks to work by @llllllllll . They could probably be further improved if people wanted to invest modest intellectual effort here. Costs scale with the number of type signatures on each operation. In blaze this was very high, well into the hundreds, in our case it would be, I think, more modest around 2-10. (also, historical note, multipledispatch predates my involvement in Blaze).

When possible it would be useful to upstream these concerns to NumPy, even if we have to move faster than NumPy is able to support.

[shoyer]

Dispatch for stack/concatenate is definitely on the radar for NumPy development, but I don't know when it's actually going to happen. The likely interface is something like __array_ufunc__: a special method like __array_concatenate__ is called on each element in the list, until one does not return NotImplemented. This is a different style of overloads than multipledispatch, one that is slightly simpler to implement but possibly slower and with fewer guarantees of correctness.

We only need this for a couple of operations, so in any case we can probably implement our own ad-hoc dispatch system for np.stack and np.concatenate, either along the of multipledispatch or NumPy/__array_ufunc__.

On further contemplation, overloading based on union types with a system like multipledispatch does seem tricky. It's not clear to me that there's even a well defined type for inputs to concatenate that should be dispatched to dask vs. numpy, for example. We want to let that dask handle any cases where at least one input is a dask array, but a type like List[Union[np.ndarray, da.Array]] actually matches a list of all numpy arrays, too -- unless we require an exact match for the type.

[llllllllll]

In blaze we have variadic sequences for multiple dispatch, and the List[Union] case is something we have run into. We have a type called VarArgs which takes a variadic sequence of type-arguments and represents a sequence of a unions over the arguments, for example: VarArgs[pd.Series, pd.DataFrame] is a sequence of unknown length which is known to hold either series or dataframes. With some mild metaprogramming we made it so that VarArs[pd.Series] is a subclass of VarArgs[pd.Series, pd.DataFrame], or in general, more specific sequences are subclasses of more general sequences. This means that you can solve the ambiguity by registering a dispatch for VarArgs[np.ndarray] and VarArgs[np.ndarray, da.Array] and you know that the second function can only be called if the sequence holds at least one dask array.

Here is an example of what that looks like for merge, which is concat(axis=1): https://github.com/blaze/blaze/blob/master/blaze/compute/pandas.py#L691
This is the definition of VarArgs: https://github.com/blaze/blaze/blob/master/blaze/compute/varargs.py

[shoyer]

@llllllllll very cool! Is there a special trick I need to use this? I tried:

# first: pip install https://github.com/blaze/blaze/archive/master.tar.gz
import blaze.compute
from blaze.compute.varargs import VarArgs
from multipledispatch import dispatch

@dispatch(VarArgs[float])
def f(args):
  print('floats')

@dispatch(VarArgs[str])
def f(args):
  print('strings')

@dispatch(VarArgs[str, float])
def f(args):
  print('mixed')

This gives me an error when I try to use it:

>>> f(['foo'])
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/multipledispatch/dispatcher.py in __call__(self, *args, **kwargs)
    154         try:
--> 155             func = self._cache[types]
    156         except KeyError:

KeyError: (<class 'list'>,)

During handling of the above exception, another exception occurred:

NotImplementedError                       Traceback (most recent call last)
<ipython-input-5-19f52a9a1dd6> in <module>()
----> 1 f(['foo'])

/usr/local/lib/python3.6/dist-packages/multipledispatch/dispatcher.py in __call__(self, *args, **kwargs)
    159                 raise NotImplementedError(
    160                         'Could not find signature for %s: <%s>' %
--> 161                         (self.name, str_signature(types)))
    162             self._cache[types] = func
    163         try:

NotImplementedError: Could not find signature for f: <list>

[llllllllll]

VarArgs itself is actually a type, so you need to create instances which wrap the list argument, for example:

In [1]: from blaze.compute.varargs import VarArgs

In [2]: from multipledispatch import dispatch

In [3]: @dispatch(VarArgs[float])
   ...: def f(args):
   ...:     print('floats')
   ...:     

In [4]: @dispatch(VarArgs[str])
   ...: def f(args):
   ...:     print('strings')
   ...:     

In [5]: @dispatch(VarArgs[str, float])
   ...: def f(args):
   ...:     print('mixed')
   ...:     

In [6]: f(VarArgs(['foo']))
strings

In [7]: f(VarArgs([1.0]))
floats

In [8]: f(VarArgs([1.0, 'foo']))
mixed

In [9]: VarArgs([1.0, 'foo'])
Out[9]: VarArgs[float, str]([1.0, 'foo'])

You could hide this behind a top-level function that wraps the input for the user, or register a dispatch for list which boxes and recurses into itself.

[hameerabbasi]

Can't some wild metaprogramming make it so that [1.0, 'foo'] itself is an instance of VarArgs[float, str] (or be converted?)

[llllllllll]

We could make a particular list an instance of a particular TypedVarArgs; however, multiple dispatch uses the type() of arguments as well as issubclass to do dispatching. Multiple dispatch depends on being able to partially order types to make dispatching more efficient. The constructor of VarArgs scans for the types of the elements and constructs an instance of a new (but memoized) subclass of VarArgs which encodes the element types so that issubclass works as expected. The problem is that type([1.0, 'foo']) returns just list which erases all information about the elements.

[llllllllll]

The wrapping dispatch would just look like:

@dispatch(list)
def f(args):
    return f(VarArgs(args))

[hameerabbasi]

How about something like checking inside a list if something is top priority, then call a, if second priority, call b, etc.

[shoyer]

Yes, I just tested out the wrapping dispatch. It works and is quite clean.

[shoyer]

As for my last concern, "Dispatch for the first argument(s) only" it looks like the simple answer is that multipledispatch already only dispatches based on positional arguments. So as long as we're strict about using keyword arguments for extra parameters like axis (which is good style anyways), we only need a single overload per array type for single dispatch functions like sum().

It looks like this resolves almost all of my concerns about using multiple dispatch.

One thing that would be nice is it VarArgs is actually distributed as part of multipledispatch rather than needing to be copied separately into xarray. That would make it easier for third parties to extend our operations, by simply importing VarArgs from multipledispatch rather than importing it from somewhere internal in xarray.

[shoyer]

How about something like checking inside a list if something is top priority, then call a, if second priority, call b, etc.

Usually, this is not a good idea. The problem is that it's impossible to know a global priority order across unrelated packages. It's usually better to declare valid type matches explicitly.

NumPy tried this with __array_priority__, but in practice these priority numbers are basically meaningless for all comparisons other than comparisons to the priority of NumPy arrays.