apply_ufunc: Add meta kwarg + bump dask to 2.2

[dcherian]

• Closes apply_ufunc with dask='parallelized' and vectorize=True fails on compute_meta #3574
• Tests added
• Passes black . && mypy . && flake8
• Fully documented, including whats-new.rst for all changes and api.rst for new API

This makes vectorize=True works with functions that add new dimensions.

cc @smartass101 @jbusecke

[dcherian]

The test failed because meta was not valid for dask==1.2

Our minimum versions policy says we can bump dask to 2.1.0.
dask 1.2 (2019-04-13) 2.1 (2019-07-08) <

So I've bumped the min dask version and added meta as a kwarg to apply_ufunc.

[dcherian]

min-all-deps fails with this weird error:

    def test_append_with_new_variable(self):
    
        ds, ds_to_append, ds_with_new_var = create_append_test_data()
    
        # check append mode for new variable
        with self.create_zarr_target() as store_target:
            xr.concat([ds, ds_to_append], dim="time").to_zarr(store_target, mode="w")
            ds_with_new_var.to_zarr(store_target, mode="a")
            combined = xr.concat([ds, ds_to_append], dim="time")
            combined["new_var"] = ds_with_new_var["new_var"]
>           assert_identical(combined, xr.open_zarr(store_target))

xarray/tests/test_backends.py:1876: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
xarray/core/dataset.py:1354: in identical
    other, "identical"
xarray/core/dataset.py:1302: in _all_compat
    self._variables, other._variables, compat=compat
xarray/core/utils.py:338: in dict_equiv
    if k not in second or not compat(first[k], second[k]):
xarray/core/dataset.py:1299: in compat
    return getattr(x, compat_str)(y)
xarray/core/variable.py:1670: in identical
    other, equiv=equiv
xarray/core/variable.py:1647: in equals
    self._data is other._data or equiv(self.data, other.data)
xarray/core/duck_array_ops.py:224: in array_equiv
    flag_array = (arr1 == arr2) | (isnull(arr1) & isnull(arr2))
/usr/share/miniconda/envs/xarray-tests/lib/python3.6/site-packages/dask/array/core.py:1739: in __eq__
    return elemwise(operator.eq, self, other)
/usr/share/miniconda/envs/xarray-tests/lib/python3.6/site-packages/dask/array/core.py:3749: in elemwise
    **blockwise_kwargs
/usr/share/miniconda/envs/xarray-tests/lib/python3.6/site-packages/dask/array/blockwise.py:145: in blockwise
    chunkss, arrays = unify_chunks(*args)
/usr/share/miniconda/envs/xarray-tests/lib/python3.6/site-packages/dask/array/core.py:3020: in unify_chunks
    (asanyarray(a) if ind is not None else a, ind) for a, ind in partition(2, args)
/usr/share/miniconda/envs/xarray-tests/lib/python3.6/site-packages/dask/array/core.py:3020: in <listcomp>
    (asanyarray(a) if ind is not None else a, ind) for a, ind in partition(2, args)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

a = array(['2019-01-01T00:00:00.000000000', '2019-01-02T00:00:00.000000000',
       '2019-01-03T00:00:00.000000000', '2019-01-04T00:00:00.000000000',
       '2019-01-05T00:00:00.000000000'], dtype='datetime64[ns]')

    def asanyarray(a):
        """Convert the input to a dask array.
    
        Subclasses of ``np.ndarray`` will be passed through as chunks unchanged.
    
        Parameters
        ----------
        a : array-like
            Input data, in any form that can be converted to a dask array.
    
        Returns
        -------
        out : dask array
            Dask array interpretation of a.
    
        Examples
        --------
        >>> import dask.array as da
        >>> import numpy as np
        >>> x = np.arange(3)
        >>> da.asanyarray(x)
        dask.array<array, shape=(3,), dtype=int64, chunksize=(3,)>
    
        >>> y = [[1, 2, 3], [4, 5, 6]]
        >>> da.asanyarray(y)
        dask.array<array, shape=(2, 3), dtype=int64, chunksize=(2, 3)>
        """
        if isinstance(a, Array):
            return a
        elif hasattr(a, "to_dask_array"):
            return a.to_dask_array()
>       elif hasattr(a, "data") and type(a).__module__.startswith("xarray."):
E       ValueError: cannot include dtype 'M' in a buffer

[dcherian]

Does anyone have thoughts on what might be happening here?

[shoyer]

This looks like an dask bug that we fixed a somewhat recently (6 months ago?).

I wouldn’t worry about it too much about it. Supporting old versions of dependencies is a best effort, not a mandate. I would just skip it on old versions of dask and move on.

[dcherian]

Ah found it: dask/dask#5334. Missed it when I looked through the changelog earlier.

[dcherian]

ready for review and merge.

[shoyer]

OK great! Looks good to me!

…

On Tue, Jan 21, 2020 at 9:32 AM Deepak Cherian ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In xarray/tests/test_sparse.py <#3660 (comment)>: > @@ -873,3 +873,16 @@ def test_dask_token(): t5 = dask.base.tokenize(ac + 1) assert t4 != t5 assert isinstance(ac.data._meta, sparse.COO) + + ***@***.***_dask There's a importorskip("sparse") on top so I think this is fine — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#3660?email_source=notifications&email_token=AAJJFVVGGV56U62XNSJLVNTQ64WURA5CNFSM4KCEN2J2YY3PNVWWK3TUL52HS4DFWFIHK3DMKJSXC5LFON2FEZLWNFSXPKTDN5WW2ZLOORPWSZGOCSQHUMA#discussion_r369142143>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVVWVJ2KZUF3NYURBM3Q64WURANCNFSM4KCEN2JQ> .

[dcherian]

Thanks for the review. I'll merge when tests pass.