From a629c3959f0451ad15b44ccc3311490b5e8cec6e Mon Sep 17 00:00:00 2001 From: inchang-ing <197932532+inchang-ing@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:24:42 +0800 Subject: [PATCH 1/3] Keep input chunks in rolling reduction results When a dask-backed array is rolled and then reduced, the rolling dimension of the result could come out with a different chunk structure than the input: padding appends its padding as separate chunks and dask's sliding_window_view rechunks so that every window chunk is complete. center=True split the dimension into an extra tiny chunk (e.g. (N-1, 1)), but center=False with an already chunked rolling dimension also came out misaligned, and the bottleneck path produced yet another chunking. Rechunk the result along the rolled dimensions back to the input chunks at the Rolling level, which covers every reduction path (the numpy, bottleneck and mean/count variants), both center settings and any chunked array type via the ChunkManager. This only changes chunking, not values. Fixes #11689 --- doc/whats-new.rst | 10 +++++++++ xarray/computation/rolling.py | 38 +++++++++++++++++++++++++++++++---- xarray/tests/test_rolling.py | 17 ++++++++++++++++ 3 files changed, 61 insertions(+), 4 deletions(-) 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..ebf219995df 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(size != 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,10 @@ 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)) From b9befdbd1961adcd4e31b4eed00580341d846069 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:26:27 +0000 Subject: [PATCH 2/3] style: pre-commit fixes --- xarray/computation/rolling.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/xarray/computation/rolling.py b/xarray/computation/rolling.py index ebf219995df..6d8677717d9 100644 --- a/xarray/computation/rolling.py +++ b/xarray/computation/rolling.py @@ -774,9 +774,7 @@ def _bottleneck_reduce(self, func, keep_attrs, **kwargs): attrs = self.obj.attrs if keep_attrs else {} return self._restore_input_chunks( - self.obj.__class__( - values, self.obj.coords, attrs=attrs, name=self.obj.name - ) + self.obj.__class__(values, self.obj.coords, attrs=attrs, name=self.obj.name) ) def _array_reduce( From 027c84e3d1ad759806c4d5741a062bc6dcd9d4e7 Mon Sep 17 00:00:00 2001 From: inchang-ing <197932532+inchang-ing@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:23:53 +0800 Subject: [PATCH 3/3] Use math.isnan for the unknown-chunk check ruff flags as PLR0124 (name compared with itself). --- xarray/computation/rolling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/computation/rolling.py b/xarray/computation/rolling.py index 6d8677717d9..0aadf915ad2 100644 --- a/xarray/computation/rolling.py +++ b/xarray/computation/rolling.py @@ -296,7 +296,7 @@ def _restore_input_chunks(self, obj: T_Xarray) -> T_Xarray: # 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(size != size for size in dim_chunks) + or any(math.isnan(size) for size in dim_chunks) ): continue chunks[dim] = dim_chunks