[inchang-ing]

Description

Fixes #11689.

When a dask-backed array is rolled and then reduced (e.g. .median()), the rolling
dimension of the result could come out with a different chunk structure than the
input. The extra tiny chunk breaks downstream operations that use the rolling
dimension as a core dimension (e.g. ufuncs).

Root cause

Variable._rolling_window pads the array before calling sliding_window_view.
dask.array.pad appends the right-hand padding as separate chunks, and dask's
sliding_window_view rechunks so that every window chunk is complete. Both change
the chunking along the rolling dimension of the result:

• center=True split the dimension into an extra tiny chunk (e.g. (N-1, 1)).
• center=False with an already chunked rolling dimension also came out
  misaligned (e.g. (10, 10) → (14, 6)).
• the bottleneck path (_bottleneck_reduce, the default for many reductions)
  produced yet another chunking (e.g. (17, 3)), because it does its own dask
  padding workaround.

Fix

Rechunk the result along the rolled dimensions back to the input chunks at the
Rolling level (Rolling._restore_input_chunks), applied to:

• Rolling.reduce (covers the numpy path),
• DataArrayRolling._bottleneck_reduce (covers the bottleneck path),
• Rolling.count (and thereby mean, which is sum / count).

This covers every reduction path, both center settings, and any chunked array
type, since it goes through the ChunkManager (obj.chunk(...)). The cost is one
extra rechunk in the graph. It only changes chunking, not values.

construct() is deliberately not touched: its chunking is controlled by the
automatic_rechunk argument of sliding_window_view (see
test_rolling_construct_automatic_rechunk), and rechunking a windowed result
back to the input chunks would be quadratic in the number of chunks.

Checklist

• User-facing changes come with a whats-new entry (doc/whats-new.rst)
• New regression test added in xarray/tests/test_rolling.py
• Full xarray/tests/test_rolling.py passes (3152 passed, 691 skipped)

Testing notes

The new test (test_rolling_dask_result_keeps_input_chunks) asserts
result.chunksizes == da.chunksizes across a matrix of input chunkings
((20,), (10, 10), (5, 5, 5, 5)), center on/off, bottleneck on/off,
numbagg on/off, and mean/sum/count/median/std/max (144 cases).
Before the fix, only the single-chunk center=True non-bottleneck case was
aligned.

[headtr1ck]

Thanks for looking into this! I think this fix is too narrow. It targets one symptom in one code path, but the misaligned chunks come up in more places than the issue shows. Here are the output chunks along s for da.rolling(s=5, center=...).mean() (input chunked with t=1 plus the listed chunks along s):

input chunks	center	bottleneck	main	this PR
(20,)	True	False	(19, 1)	(20,) ✅
(20,)	True	True	(17, 3)	(17, 3) ❌
(10, 10)	False	False	(14, 6)	(14, 6) ❌
(10, 10)	True	False	(12, 7, 1)	(12, 8) ❌
(10, 10)	True	True	(8, 9, 3)	(8, 9, 3) ❌
(5, 5, 5, 5)	False	False	(9, 5, 5, 1)	(9, 5, 5, 1) ❌
(5, 5, 5, 5)	True	True	(3, 5, 5, 7)	(3, 5, 5, 7) ❌

import numpy as np
import xarray as xr

da = xr.DataArray(np.ones((4, 20)), dims=("t", "s")).chunk(t=1, s=10)
with xr.set_options(use_bottleneck=False, use_numbagg=False):
    print(da.rolling(s=5).mean().chunksizes["s"])  # (14, 6)

So:

• The default bottleneck path (_bottleneck_reduce, which has its own dask padding workaround) is not covered. The new test disables bottleneck/numbagg, which hides that.
• center=True is not really the cause. With center=False, an input that is already chunked along the rolling dimension also comes out with different chunks.
• The fix depends on how dask.pad and sliding_window_view happen to chunk the padding. It is also dask-only (is_duck_dask_array and .rechunk), so other chunked array types are not covered.

A more general approach might be to rechunk the result along the rolled dimensions back to the input chunks at the Rolling level, through the ChunkManager (e.g. Variable.chunk), so it covers every reduction path, both center settings and any chunked array type. The test could then simply assert result.chunksizes == da.chunksizes with bottleneck/numbagg on and off. The cost is an extra rechunk in the graph, so it would be good to hear from others whether that's acceptable, or whether this should rather be fixed upstream in dask's sliding_window_view(automatic_rechunk=True).

cc @dcherian @max-sixty, who have worked on the dask/numbagg rolling paths.

[This is Claude Code on behalf of Michael Niklas]

[inchang-ing]

Thanks for the thorough review — you're right that the first attempt was too narrow, and your matrix was exactly what I needed. I've reworked the PR along the lines you suggested:

New approach (a629c395): rechunk the result along the rolled dimensions back to the input chunks at the Rolling level (Rolling._restore_input_chunks, going through the ChunkManager via obj.chunk(...)), applied to

• Rolling.reduce (numpy path),
• DataArrayRolling._bottleneck_reduce (bottleneck path — the default for many reductions, which your matrix showed was not covered),
• Rolling.count, which also fixes mean (sum / count previously unified the aligned sum chunks with the misaligned count chunks).

This covers every reduction path, both center settings and any chunked array type, as you suggested. The old variable.py pad-merging fix is gone. I verified your matrix: all seven rows now produce the input chunks.

On construct(): I deliberately left it untouched. Its chunking is controlled by the sliding_window_view_kwargs={'automatic_rechunk': ...} argument (test_rolling_construct_automatic_rechunk asserts the outputs differ), and wrapping it also makes dask's rechunk planner blow up: with 400 chunks along time and window=100, rechunking the windowed result back to the input chunks hangs in dask.array.rechunk._compute_rechunk (the chunk count along the rolling dimension grows by the window size). Since the issue is about reductions, I'd rather keep the construct semantics intact — happy to hear if you disagree.

Test: test_rolling_dask_result_keeps_input_chunks now asserts result.chunksizes == da.chunksizes over a matrix of input chunkings × center × bottleneck × numbagg × {mean, sum, count, median, std, max} (144 cases). Full xarray/tests/test_rolling.py: 3152 passed, 691 skipped.

On the open question you raised (fix here vs. upstream in dask's sliding_window_view): I agree the dask behavior is arguably the root cause, but the extra rechunk here is cheap (same total size, no data movement across the network for a single-machine scheduler) and keeps xarray's output predictable regardless of dask version. If the maintainers prefer an upstream fix, this PR can be reduced to a workaround note.