diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 8f7566e83fd..fc363e96b0f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -363,6 +363,16 @@ Performance By `Michael Niklas `_. +Bug fixes +~~~~~~~~~ +- Rolling reductions of dask-backed arrays now keep the input chunks along the + rolled dimensions in the result. Previously, padding and building the sliding + windows could split the rolling dimension into differently-sized chunks + (e.g. ``(N-1, 1)`` for ``center=True``), which broke downstream operations + that use the rolling dimension as a core dimension (:issue:`11689`). + By `inchang-ing `_. + + Internal Changes ~~~~~~~~~~~~~~~~ - The generated arithmetic operators in ``xarray/core/_typed_ops.py`` use PEP 695 diff --git a/xarray/computation/rolling.py b/xarray/computation/rolling.py index 90e5087aaf8..0aadf915ad2 100644 --- a/xarray/computation/rolling.py +++ b/xarray/computation/rolling.py @@ -272,10 +272,38 @@ def count(self, keep_attrs: bool | None = None) -> T_Xarray: keep_attrs = self._get_keep_attrs(keep_attrs) rolling_count = self._counts(keep_attrs=keep_attrs) enough_periods = rolling_count >= self.min_periods - return rolling_count.where(enough_periods) + return self._restore_input_chunks(rolling_count.where(enough_periods)) count.__doc__ = _ROLLING_REDUCE_DOCSTRING_TEMPLATE.format(name="count") + def _restore_input_chunks(self, obj: T_Xarray) -> T_Xarray: + """Rechunk the rolled dimensions back to the input chunks. + + Padding the array and building the sliding windows changes the chunk + structure along the rolled dimensions: ``pad`` appends its padding as + separate chunks, and dask's ``sliding_window_view`` rechunks so that + every window chunk is complete. Rechunking the result back to the + input chunks keeps the output chunk structure predictable, which + downstream operations that use a rolled dimension as a core dimension + rely on (GH #11689). + """ + obj_chunksizes = self.obj.chunksizes # empty when not chunked + chunks = {} + for dim in self.dim: + dim_chunks = obj_chunksizes.get(dim) + if ( + dim_chunks is None + # a strided construct changes the length of the dimension + or obj.sizes[dim] != self.obj.sizes[dim] + # NaN marks unknown chunks, which cannot be rechunked to + or any(math.isnan(size) for size in dim_chunks) + ): + continue + chunks[dim] = dim_chunks + if not chunks: + return obj + return obj.chunk(chunks) + def _mapping_to_list( self, arg: _T | Mapping[Any, _T], @@ -623,7 +651,7 @@ def reduce( # Find valid windows based on count. counts = self._counts(keep_attrs=False) - return result.where(counts >= self.min_periods) + return self._restore_input_chunks(result.where(counts >= self.min_periods)) @override def _counts(self, keep_attrs: bool | None) -> DataArray: @@ -745,8 +773,8 @@ def _bottleneck_reduce(self, func, keep_attrs, **kwargs): attrs = self.obj.attrs if keep_attrs else {} - return self.obj.__class__( - values, self.obj.coords, attrs=attrs, name=self.obj.name + return self._restore_input_chunks( + self.obj.__class__(values, self.obj.coords, attrs=attrs, name=self.obj.name) ) def _array_reduce( diff --git a/xarray/tests/test_rolling.py b/xarray/tests/test_rolling.py index d96ebf151fe..4411eb93d02 100644 --- a/xarray/tests/test_rolling.py +++ b/xarray/tests/test_rolling.py @@ -196,6 +196,23 @@ def test_rolling_dask_keeps_chunks(self, name) -> None: assert actual.chunksizes["y"] == da.chunksizes["y"] assert len(actual.chunksizes["x"]) == len(da.chunksizes["x"]) + @requires_dask + @pytest.mark.parametrize("chunks", ((20,), (10, 10), (5, 5, 5, 5))) + @pytest.mark.parametrize("center", (True, False)) + @pytest.mark.parametrize("use_bottleneck", (True, False)) + @pytest.mark.parametrize("use_numbagg", (True, False)) + @pytest.mark.parametrize("name", ("mean", "sum", "count", "median", "std", "max")) + def test_rolling_dask_result_keeps_input_chunks( + self, name, use_bottleneck, use_numbagg, center, chunks + ) -> None: + # The result must keep the input chunks along the rolled dimension on + # every reduction path (bottleneck/numbagg on and off, center on and + # off): https://github.com/pydata/xarray/issues/11689 + da = DataArray(np.ones((4, 20)), dims=("t", "s")).chunk(t=1, s=chunks) + with set_options(use_bottleneck=use_bottleneck, use_numbagg=use_numbagg): + actual = getattr(da.rolling(s=5, center=center), name)() + assert actual.chunksizes == da.chunksizes + @requires_dask @pytest.mark.parametrize("name", ("mean", "count")) @pytest.mark.parametrize("center", (True, False, None))