Enable taking the mean of dask-backed cftime arrays

[spencerkclark]

This was essentially enabled by @dcherian in #6556, but we did not remove the error that prevented computing the mean of a dask-backed cftime array. This PR removes that error, and adds some tests. One minor modification in _timedelta_to_seconds was needed for compatibility with scalar cftime arrays.

This happens to address the second part of #5897, so I added a regression test for that. It seems like we decided to simply document the behavior in the first part (#5898, dcherian@99bfe12), but I'm not sure if we intend to change that behavior eventually or not.

• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

[spencerkclark]

Ah, this still didn't quite work for older dask versions:

Details

====================================================================================================== FAILURES =======================================================================================================
___________________________________________________________________________________________ test_cftime_datetime_mean[True] ___________________________________________________________________________________________

dask = True

    @requires_cftime
    @pytest.mark.parametrize("dask", [False, True])
    def test_cftime_datetime_mean(dask):
        if dask and not has_dask:
            pytest.skip("requires dask")

        times = cftime_range("2000", periods=4)
        da = DataArray(times, dims=["time"])
        da_2d = DataArray(times.values.reshape(2, 2))

        if dask:
            da = da.chunk({"time": 2})
            da_2d = da_2d.chunk({"dim_0": 2})

        expected = da.isel(time=0)
        # one compute needed to check the array contains cftime datetimes
        with raise_if_dask_computes(max_computes=1):
>           result = da.isel(time=0).mean()

/Users/spencer/software/xarray/xarray/tests/test_duck_array_ops.py:342:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/Users/spencer/software/xarray/xarray/core/_reductions.py:1478: in mean
    return self.reduce(
/Users/spencer/software/xarray/xarray/core/dataarray.py:2930: in reduce
    var = self.variable.reduce(func, dim, axis, keep_attrs, keepdims, **kwargs)
/Users/spencer/software/xarray/xarray/core/variable.py:1854: in reduce
    data = func(self.data, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

array = dask.array<getitem, shape=(), dtype=object, chunksize=(), chunktype=numpy.ndarray>, axis = None, skipna = None, kwargs = {}
_contains_cftime_datetimes = <function _contains_cftime_datetimes at 0x7fcdd91ce3a0>, offset = dask.array<where, shape=(), dtype=object, chunksize=(), chunktype=numpy.ndarray>
timedeltas = dask.array<multiply, shape=(), dtype=float64, chunksize=(), chunktype=numpy.ndarray>, mean_timedeltas = dask.array<mean_agg-aggregate, shape=(), dtype=float64, chunksize=(), chunktype=numpy.ndarray>

    def mean(array, axis=None, skipna=None, **kwargs):
        """inhouse mean that can handle np.datetime64 or cftime.datetime
        dtypes"""
        from .common import _contains_cftime_datetimes

        array = asarray(array)
        if array.dtype.kind in "Mm":
            offset = _datetime_nanmin(array)

            # xarray always uses np.datetime64[ns] for np.datetime64 data
            dtype = "timedelta64[ns]"
            return (
                _mean(
                    datetime_to_numeric(array, offset), axis=axis, skipna=skipna, **kwargs
                ).astype(dtype)
                + offset
            )
        elif _contains_cftime_datetimes(array):
            offset = min(array)
            timedeltas = datetime_to_numeric(array, offset, datetime_unit="us")
            mean_timedeltas = _mean(timedeltas, axis=axis, skipna=skipna, **kwargs)
>           return _to_pytimedelta(mean_timedeltas, unit="us") + offset
E           TypeError: unsupported operand type(s) for +: 'Array' and 'Array'

/Users/spencer/software/xarray/xarray/core/duck_array_ops.py:573: TypeError

This appears to be the commit that fixed this issue in dask, meaning it should be fixed as of dask version 2021.07.0:

git bisect broken
626cc724ffa1d14cd74f48ff4f21ef443d832afa is the first fixed commit
commit 626cc724ffa1d14cd74f48ff4f21ef443d832afa
Author: Julia Signell <jsignell@gmail.com>
Date:   Tue Jun 29 11:50:36 2021 -0400

    Allow mixing dask and numpy arrays in @guvectorize (#6863)

 dask/array/core.py              |  7 ++++++-
 dask/array/tests/test_gufunc.py | 11 +++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

[dcherian]

Thanks @spencerkclark

dask-core         2021.4  (2021-04-02) 2021.8  (2021-08-13) <
distributed       2021.4  (2021-04-02) 2021.8  (2021-08-14) <

It looks like we can bump to 2021.08.0 instead of dealing with this compatibility code.

[spencerkclark]

It looks like we can bump to 2021.08.0 instead of dealing with this compatibility code.

Thanks for noting this @dcherian. I wasn't sure if this was worth bumping our minimum version for, but since our policy allows for it, I'm happy to make that change.

[dcherian]

nice! Thanks for making extra sure that we didn't compute the whole thing.

[dcherian]

LGTM!

[mathause]

Looks good - can you merge master to fix the merge conflict?

[spencerkclark]

Oops sorry I missed the merge conflict. Thanks @dcherian and @mathause.