diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 2f1d76d7abc..ab45f4773e3 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -57,6 +57,14 @@ Internal Changes runtime behavior. This enables CI integration for type stub validation and helps prevent type annotation regressions (:issue:`11086`). By `Kristian KollsgÄrd `_. +- Add flox support for :py:meth:`DataArray.groupby().median`, + :py:meth:`Dataset.groupby().median`, :py:meth:`DataArray.resample().median`, and + :py:meth:`Dataset.resample().median`. This significantly speeds up median reductions + when flox is installed and the data chunking allows blockwise processing. For + incompatible chunking, a fallback to the non-flox implementation ensures backward + compatibility. (:issue:`11238`, :pull:`11239`). By `Samuel Le Meur-Diebolt + `_. + .. _whats-new.2026.02.0: diff --git a/xarray/core/_aggregations.py b/xarray/core/_aggregations.py index 52f6557a8ad..fd47b8121a8 100644 --- a/xarray/core/_aggregations.py +++ b/xarray/core/_aggregations.py @@ -4911,6 +4911,28 @@ def median( Data variables: da (labels) float64 24B nan 2.0 1.5 """ + if ( + flox_available + and OPTIONS["use_flox"] + and contains_only_chunked_or_numpy(self._obj) + ): + try: + return self._flox_reduce( + func="median", + dim=dim, + skipna=skipna, + numeric_only=True, + # fill_value=fill_value, + keep_attrs=keep_attrs, + **kwargs, + ) + except ValueError as e: + if "median is only supported for `method='blockwise'`" in str(e): + # Fall back to non-flox implementation when chunking doesn't + # allow blockwise processing + pass + else: + raise return self.reduce( duck_array_ops.median, dim=dim, @@ -6407,6 +6429,28 @@ def median( Data variables: da (time) float64 24B 1.0 2.0 nan """ + if ( + flox_available + and OPTIONS["use_flox"] + and contains_only_chunked_or_numpy(self._obj) + ): + try: + return self._flox_reduce( + func="median", + dim=dim, + skipna=skipna, + numeric_only=True, + # fill_value=fill_value, + keep_attrs=keep_attrs, + **kwargs, + ) + except ValueError as e: + if "median is only supported for `method='blockwise'`" in str(e): + # Fall back to non-flox implementation when chunking doesn't + # allow blockwise processing + pass + else: + raise return self.reduce( duck_array_ops.median, dim=dim, @@ -7804,6 +7848,27 @@ def median( Coordinates: * labels (labels) object 24B 'a' 'b' 'c' """ + if ( + flox_available + and OPTIONS["use_flox"] + and contains_only_chunked_or_numpy(self._obj) + ): + try: + return self._flox_reduce( + func="median", + dim=dim, + skipna=skipna, + # fill_value=fill_value, + keep_attrs=keep_attrs, + **kwargs, + ) + except ValueError as e: + if "median is only supported for `method='blockwise'`" in str(e): + # Fall back to non-flox implementation when chunking doesn't + # allow blockwise processing + pass + else: + raise return self.reduce( duck_array_ops.median, dim=dim, @@ -9192,6 +9257,27 @@ def median( Coordinates: * time (time) datetime64[us] 24B 2001-01-31 2001-04-30 2001-07-31 """ + if ( + flox_available + and OPTIONS["use_flox"] + and contains_only_chunked_or_numpy(self._obj) + ): + try: + return self._flox_reduce( + func="median", + dim=dim, + skipna=skipna, + # fill_value=fill_value, + keep_attrs=keep_attrs, + **kwargs, + ) + except ValueError as e: + if "median is only supported for `method='blockwise'`" in str(e): + # Fall back to non-flox implementation when chunking doesn't allow + # blockwise processing. + pass + else: + raise return self.reduce( duck_array_ops.median, dim=dim,