diff --git a/doc/whats-new.rst b/doc/whats-new.rst index dc4fb7ae722..cbc74cdb8c4 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -38,6 +38,8 @@ New Features for the calculation of nanquantiles (i.e., `skipna=True`) if it is installed. This is currently limited to the linear interpolation method (`method='linear'`). (:issue:`7377`, :pull:`8684`) By `Marco Wolsza `_. +- Grouped and resampled median and quantile calculations now use ``flox>=0.9.1`` if present. + By `Deepak Cherian `_. Breaking changes ~~~~~~~~~~~~~~~~ diff --git a/xarray/core/_aggregations.py b/xarray/core/_aggregations.py index bee6afd5a19..12e95f281af 100644 --- a/xarray/core/_aggregations.py +++ b/xarray/core/_aggregations.py @@ -2392,8 +2392,6 @@ def count( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -2490,8 +2488,6 @@ def all( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -2588,8 +2584,6 @@ def any( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -2692,8 +2686,6 @@ def max( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -2808,8 +2800,6 @@ def min( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -2924,8 +2914,6 @@ def mean( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -3049,8 +3037,6 @@ def prod( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -3186,8 +3172,6 @@ def sum( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -3320,8 +3304,6 @@ def std( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -3454,8 +3436,6 @@ def var( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -3584,8 +3564,6 @@ def median( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -3628,14 +3606,30 @@ def median( Data variables: da (labels) float64 24B nan 2.0 1.5 """ - return self._reduce_without_squeeze_warn( - duck_array_ops.median, - dim=dim, - skipna=skipna, - numeric_only=True, - keep_attrs=keep_attrs, - **kwargs, - ) + if ( + flox_available + and OPTIONS["use_flox"] + and module_available("flox", minversion="0.9.1") + and contains_only_chunked_or_numpy(self._obj) + ): + return self._flox_reduce( + func="median", + dim=dim, + skipna=skipna, + numeric_only=True, + # fill_value=fill_value, + keep_attrs=keep_attrs, + **kwargs, + ) + else: + return self._reduce_without_squeeze_warn( + duck_array_ops.median, + dim=dim, + skipna=skipna, + numeric_only=True, + keep_attrs=keep_attrs, + **kwargs, + ) def cumsum( self, @@ -3687,8 +3681,6 @@ def cumsum( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -3788,8 +3780,6 @@ def cumprod( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -3919,8 +3909,6 @@ def count( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -4017,8 +4005,6 @@ def all( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -4115,8 +4101,6 @@ def any( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -4219,8 +4203,6 @@ def max( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -4335,8 +4317,6 @@ def min( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -4451,8 +4431,6 @@ def mean( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -4576,8 +4554,6 @@ def prod( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -4713,8 +4689,6 @@ def sum( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -4847,8 +4821,6 @@ def std( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -4981,8 +4953,6 @@ def var( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -5111,8 +5081,6 @@ def median( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -5155,14 +5123,30 @@ def median( Data variables: da (time) float64 24B 1.0 2.0 nan """ - return self._reduce_without_squeeze_warn( - duck_array_ops.median, - dim=dim, - skipna=skipna, - numeric_only=True, - keep_attrs=keep_attrs, - **kwargs, - ) + if ( + flox_available + and OPTIONS["use_flox"] + and module_available("flox", minversion="0.9.1") + and contains_only_chunked_or_numpy(self._obj) + ): + return self._flox_reduce( + func="median", + dim=dim, + skipna=skipna, + numeric_only=True, + # fill_value=fill_value, + keep_attrs=keep_attrs, + **kwargs, + ) + else: + return self._reduce_without_squeeze_warn( + duck_array_ops.median, + dim=dim, + skipna=skipna, + numeric_only=True, + keep_attrs=keep_attrs, + **kwargs, + ) def cumsum( self, @@ -5214,8 +5198,6 @@ def cumsum( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -5315,8 +5297,6 @@ def cumprod( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -5446,8 +5426,6 @@ def count( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -5537,8 +5515,6 @@ def all( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -5628,8 +5604,6 @@ def any( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -5725,8 +5699,6 @@ def max( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -5832,8 +5804,6 @@ def min( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Examples @@ -5939,8 +5909,6 @@ def mean( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -6055,8 +6023,6 @@ def prod( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -6181,8 +6147,6 @@ def sum( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -6304,8 +6268,6 @@ def std( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -6427,8 +6389,6 @@ def var( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -6546,8 +6506,6 @@ def median( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -6583,13 +6541,28 @@ def median( Coordinates: * labels (labels) object 24B 'a' 'b' 'c' """ - return self._reduce_without_squeeze_warn( - duck_array_ops.median, - dim=dim, - skipna=skipna, - keep_attrs=keep_attrs, - **kwargs, - ) + if ( + flox_available + and OPTIONS["use_flox"] + and module_available("flox", minversion="0.9.1") + and contains_only_chunked_or_numpy(self._obj) + ): + return self._flox_reduce( + func="median", + dim=dim, + skipna=skipna, + # fill_value=fill_value, + keep_attrs=keep_attrs, + **kwargs, + ) + else: + return self._reduce_without_squeeze_warn( + duck_array_ops.median, + dim=dim, + skipna=skipna, + keep_attrs=keep_attrs, + **kwargs, + ) def cumsum( self, @@ -6641,8 +6614,6 @@ def cumsum( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -6738,8 +6709,6 @@ def cumprod( Use the ``flox`` package to significantly speed up groupby computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - other methods might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -6865,8 +6834,6 @@ def count( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -6956,8 +6923,6 @@ def all( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -7047,8 +7012,6 @@ def any( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -7144,8 +7107,6 @@ def max( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -7251,8 +7212,6 @@ def min( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Examples @@ -7358,8 +7317,6 @@ def mean( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -7474,8 +7431,6 @@ def prod( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -7600,8 +7555,6 @@ def sum( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -7723,8 +7676,6 @@ def std( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -7846,8 +7797,6 @@ def var( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -7965,8 +7914,6 @@ def median( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -8002,13 +7949,28 @@ def median( Coordinates: * time (time) datetime64[ns] 24B 2001-01-31 2001-04-30 2001-07-31 """ - return self._reduce_without_squeeze_warn( - duck_array_ops.median, - dim=dim, - skipna=skipna, - keep_attrs=keep_attrs, - **kwargs, - ) + if ( + flox_available + and OPTIONS["use_flox"] + and module_available("flox", minversion="0.9.1") + and contains_only_chunked_or_numpy(self._obj) + ): + return self._flox_reduce( + func="median", + dim=dim, + skipna=skipna, + # fill_value=fill_value, + keep_attrs=keep_attrs, + **kwargs, + ) + else: + return self._reduce_without_squeeze_warn( + duck_array_ops.median, + dim=dim, + skipna=skipna, + keep_attrs=keep_attrs, + **kwargs, + ) def cumsum( self, @@ -8060,8 +8022,6 @@ def cumsum( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. @@ -8157,8 +8117,6 @@ def cumprod( Use the ``flox`` package to significantly speed up resampling computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. - The default choice is ``method="cohorts"`` which generalizes the best, - ``method="blockwise"`` might work better for your problem. See the `flox documentation `_ for more. Non-numeric variables will be removed prior to reducing. diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 3aabf618a20..e99308ebddc 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -26,15 +26,17 @@ filter_indexes_from_coords, safe_cast_to_index, ) -from xarray.core.options import _get_keep_attrs +from xarray.core.options import OPTIONS, _get_keep_attrs from xarray.core.types import Dims, QuantileMethods, T_DataArray, T_Xarray from xarray.core.utils import ( FrozenMappingWarningOnValuesAccess, + contains_only_chunked_or_numpy, either_dict_or_kwargs, emit_user_level_warning, hashable, is_scalar, maybe_wrap_array, + module_available, peek_at, ) from xarray.core.variable import IndexVariable, Variable @@ -1267,16 +1269,26 @@ def quantile( (grouper,) = self.groupers dim = grouper.group1d.dims - return self.map( - self._obj.__class__.quantile, - shortcut=False, - q=q, - dim=dim, - method=method, - keep_attrs=keep_attrs, - skipna=skipna, - interpolation=interpolation, - ) + if ( + method == "linear" + and OPTIONS["use_flox"] + and contains_only_chunked_or_numpy(self._obj) + and module_available("flox", minversion="0.9.1") + ): + return self._flox_reduce( + func="quantile", q=q, dim=dim, keep_attrs=keep_attrs, skipna=skipna + ) + else: + return self.map( + self._obj.__class__.quantile, + shortcut=False, + q=q, + dim=dim, + method=method, + keep_attrs=keep_attrs, + skipna=skipna, + interpolation=interpolation, + ) def where(self, cond, other=dtypes.NA) -> T_Xarray: """Return elements from `self` or `other` depending on `cond`. diff --git a/xarray/util/generate_aggregations.py b/xarray/util/generate_aggregations.py index 3462af28663..d2ae1655ad2 100644 --- a/xarray/util/generate_aggregations.py +++ b/xarray/util/generate_aggregations.py @@ -19,6 +19,7 @@ MODULE_PREAMBLE = '''\ """Mixin classes with reduction operations.""" + # This file was generated using xarray.util.generate_aggregations. Do not edit manually. from __future__ import annotations @@ -245,13 +246,9 @@ def {method}( _FLOX_NOTES_TEMPLATE = """Use the ``flox`` package to significantly speed up {kind} computations, especially with dask arrays. Xarray will use flox by default if installed. Pass flox-specific keyword arguments in ``**kwargs``. -The default choice is ``method="cohorts"`` which generalizes the best, -{recco} might work better for your problem. See the `flox documentation `_ for more.""" -_FLOX_GROUPBY_NOTES = _FLOX_NOTES_TEMPLATE.format(kind="groupby", recco="other methods") -_FLOX_RESAMPLE_NOTES = _FLOX_NOTES_TEMPLATE.format( - kind="resampling", recco='``method="blockwise"``' -) +_FLOX_GROUPBY_NOTES = _FLOX_NOTES_TEMPLATE.format(kind="groupby") +_FLOX_RESAMPLE_NOTES = _FLOX_NOTES_TEMPLATE.format(kind="resampling") ExtraKwarg = collections.namedtuple("ExtraKwarg", "docs kwarg call example") skipna = ExtraKwarg( @@ -300,11 +297,13 @@ def __init__( extra_kwargs=tuple(), numeric_only=False, see_also_modules=("numpy", "dask.array"), + min_flox_version=None, ): self.name = name self.extra_kwargs = extra_kwargs self.numeric_only = numeric_only self.see_also_modules = see_also_modules + self.min_flox_version = min_flox_version if bool_reduce: self.array_method = f"array_{name}" self.np_example_array = """ @@ -445,7 +444,7 @@ def generate_code(self, method, has_keep_attrs): # numpy_groupies & flox do not support median # https://github.com/ml31415/numpy-groupies/issues/43 - method_is_not_flox_supported = method.name in ("median", "cumsum", "cumprod") + method_is_not_flox_supported = method.name in ("cumsum", "cumprod") if method_is_not_flox_supported: indent = 12 else: @@ -466,10 +465,18 @@ def generate_code(self, method, has_keep_attrs): )""" else: - return f"""\ + return ( + """\ if ( flox_available - and OPTIONS["use_flox"] + and OPTIONS["use_flox"]""" + + ( + f""" + and module_available("flox", minversion="{method.min_flox_version}")""" + if method.min_flox_version is not None + else "" + ) + + f""" and contains_only_chunked_or_numpy(self._obj) ): return self._flox_reduce( @@ -486,6 +493,7 @@ def generate_code(self, method, has_keep_attrs): keep_attrs=keep_attrs, **kwargs, )""" + ) class GenericAggregationGenerator(AggregationGenerator): @@ -522,7 +530,9 @@ def generate_code(self, method, has_keep_attrs): Method("sum", extra_kwargs=(skipna, min_count), numeric_only=True), Method("std", extra_kwargs=(skipna, ddof), numeric_only=True), Method("var", extra_kwargs=(skipna, ddof), numeric_only=True), - Method("median", extra_kwargs=(skipna,), numeric_only=True), + Method( + "median", extra_kwargs=(skipna,), numeric_only=True, min_flox_version="0.9.1" + ), # Cumulatives: Method("cumsum", extra_kwargs=(skipna,), numeric_only=True), Method("cumprod", extra_kwargs=(skipna,), numeric_only=True),