Fix rolling(center=True) splitting the rolling dimension into two chunks

doc/whats-new.rst

@@ -363,6 +363,16 @@ Performance
   By `Michael Niklas <https://github.com/headtr1ck>`_.
 
 
+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 <https://github.com/inchang-ing>`_.
+
+
 Internal Changes
 ~~~~~~~~~~~~~~~~
 - The generated arithmetic operators in ``xarray/core/_typed_ops.py`` use PEP 695

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(

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))