From 83087c6b0cb0cfa204e6da8f6245fb97fa0f710f Mon Sep 17 00:00:00 2001 From: Samuel Le Meur-Diebolt Date: Tue, 17 Mar 2026 10:35:51 +0000 Subject: [PATCH 1/2] Use flox for median in groupby/resample operations Add flox support to median() methods in: - DataArrayGroupByAggregations - DatasetGroupByAggregations - DataArrayResampleAggregations - DatasetResampleAggregations This aligns the implementation with the documentation which already claimed flox was used when available. The fix provides significant performance improvements when flox can process the data. A fallback to the non-flox implementation is included for cases where flox's median aggregation requires blockwise processing but the data chunking doesn't support it. Closes #11238 --- doc/whats-new.rst | 8 ++++ xarray/core/_aggregations.py | 86 ++++++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 2f1d76d7abc..2d6ae6f2fcb 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:`TODO`). 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, From f5bb6f5ebba999b93e2a3878542292ae976ddfda Mon Sep 17 00:00:00 2001 From: Samuel Le Meur-Diebolt Date: Tue, 17 Mar 2026 10:49:04 +0000 Subject: [PATCH 2/2] docs: update PR ID --- doc/whats-new.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 2d6ae6f2fcb..ab45f4773e3 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -62,7 +62,7 @@ Internal Changes :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:`TODO`). By `Samuel Le Meur-Diebolt + compatibility. (:issue:`11238`, :pull:`11239`). By `Samuel Le Meur-Diebolt `_.