diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 4fd9073b50e..e88337ba946 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -42,6 +42,9 @@ Documentation Internal Changes ~~~~~~~~~~~~~~~~ +- Allow chunked non-dask arrays (i.e. Cubed arrays) in groupby operations. (:pull:`7941`) + By `Tom Nicholas `_. + .. _whats-new.2023.06.0: diff --git a/xarray/core/_aggregations.py b/xarray/core/_aggregations.py index 3051502beba..4f6dce1a04a 100644 --- a/xarray/core/_aggregations.py +++ b/xarray/core/_aggregations.py @@ -9,7 +9,7 @@ from xarray.core import duck_array_ops from xarray.core.options import OPTIONS from xarray.core.types import Dims -from xarray.core.utils import contains_only_dask_or_numpy, module_available +from xarray.core.utils import contains_only_chunked_or_numpy, module_available if TYPE_CHECKING: from xarray.core.dataarray import DataArray @@ -2413,7 +2413,7 @@ def count( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="count", @@ -2511,7 +2511,7 @@ def all( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="all", @@ -2609,7 +2609,7 @@ def any( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="any", @@ -2723,7 +2723,7 @@ def max( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="max", @@ -2839,7 +2839,7 @@ def min( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="min", @@ -2957,7 +2957,7 @@ def mean( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="mean", @@ -3092,7 +3092,7 @@ def prod( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="prod", @@ -3229,7 +3229,7 @@ def sum( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="sum", @@ -3363,7 +3363,7 @@ def std( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="std", @@ -3497,7 +3497,7 @@ def var( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="var", @@ -3927,7 +3927,7 @@ def count( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="count", @@ -4025,7 +4025,7 @@ def all( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="all", @@ -4123,7 +4123,7 @@ def any( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="any", @@ -4237,7 +4237,7 @@ def max( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="max", @@ -4353,7 +4353,7 @@ def min( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="min", @@ -4471,7 +4471,7 @@ def mean( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="mean", @@ -4606,7 +4606,7 @@ def prod( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="prod", @@ -4743,7 +4743,7 @@ def sum( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="sum", @@ -4877,7 +4877,7 @@ def std( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="std", @@ -5011,7 +5011,7 @@ def var( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="var", @@ -5436,7 +5436,7 @@ def count( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="count", @@ -5527,7 +5527,7 @@ def all( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="all", @@ -5618,7 +5618,7 @@ def any( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="any", @@ -5723,7 +5723,7 @@ def max( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="max", @@ -5830,7 +5830,7 @@ def min( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="min", @@ -5939,7 +5939,7 @@ def mean( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="mean", @@ -6063,7 +6063,7 @@ def prod( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="prod", @@ -6189,7 +6189,7 @@ def sum( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="sum", @@ -6312,7 +6312,7 @@ def std( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="std", @@ -6435,7 +6435,7 @@ def var( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="var", @@ -6842,7 +6842,7 @@ def count( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="count", @@ -6933,7 +6933,7 @@ def all( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="all", @@ -7024,7 +7024,7 @@ def any( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="any", @@ -7129,7 +7129,7 @@ def max( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="max", @@ -7236,7 +7236,7 @@ def min( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="min", @@ -7345,7 +7345,7 @@ def mean( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="mean", @@ -7469,7 +7469,7 @@ def prod( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="prod", @@ -7595,7 +7595,7 @@ def sum( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="sum", @@ -7718,7 +7718,7 @@ def std( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="std", @@ -7841,7 +7841,7 @@ def var( if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="var", diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 7270a8f587f..bd0ca57f33c 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -1124,19 +1124,19 @@ def iterate_nested(nested_list): yield item -def contains_only_dask_or_numpy(obj) -> bool: - """Returns True if xarray object contains only numpy or dask arrays. +def contains_only_chunked_or_numpy(obj) -> bool: + """Returns True if xarray object contains only numpy arrays or chunked arrays (i.e. pure dask or cubed). Expects obj to be Dataset or DataArray""" from xarray.core.dataarray import DataArray - from xarray.core.pycompat import is_duck_dask_array + from xarray.core.pycompat import is_chunked_array if isinstance(obj, DataArray): obj = obj._to_temp_dataset() return all( [ - isinstance(var.data, np.ndarray) or is_duck_dask_array(var.data) + isinstance(var.data, np.ndarray) or is_chunked_array(var.data) for var in obj.variables.values() ] ) diff --git a/xarray/util/generate_aggregations.py b/xarray/util/generate_aggregations.py index efc69c46947..6c50474030d 100644 --- a/xarray/util/generate_aggregations.py +++ b/xarray/util/generate_aggregations.py @@ -22,12 +22,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Sequence +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Callable from xarray.core import duck_array_ops from xarray.core.options import OPTIONS from xarray.core.types import Dims -from xarray.core.utils import contains_only_dask_or_numpy, module_available +from xarray.core.utils import contains_only_chunked_or_numpy, module_available if TYPE_CHECKING: from xarray.core.dataarray import DataArray @@ -394,7 +395,7 @@ def generate_code(self, method): if ( flox_available and OPTIONS["use_flox"] - and contains_only_dask_or_numpy(self._obj) + and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( func="{method.name}",