Repository: dask/dask
Issue category: Performance Issue

Issue:
xarray.rolling().construct() and sliding_window_view chunking behaviour
Using ``.rolling().construct()`` will inject new dimensions into an array with ``sliding_window_view`` under the hood. NumPy is clever about this by arranging things to be a view of the underlying array and thus not increasing the memory footprint. Problems are starting to show up if you call a follow up operation that requires a copy, std is an example.

It uses construct under the hood but the aggregation requires a copy to handle NaNs, which makes memory blow up. A better example is something like

```
arr = xr.DataArray(da.random.random((2000, 2000, 2000), chunks=(200, 200, 200)), dims=['x','y', 'z'])
arr = arr.rolling({"x": 10, "y": 20}).construct({"x": "dim1", "y": "dim2"})
```

construct increase the chunk size by a factor of 200 to 14GB , things are still fine because we haven't performed a copy yet. But almost every subsequent operation will require copying the data, so things will go south real quick

```
arr.isel(dim1=[1, 3, 5, 7, 9])
```

This reduces chunk size to 7GB, but we will make a copy and memory will blow up

How to deal with this?

My suggestions would be to split up the other chunks in ``sliding_window_view`` to keep chunk sizes the same after construct, this would slow us down for cases where we indeed don't make a copy, but I think this is sufficiently rare in the general construct case.

Construct is also used internally for reductions, splitting up chunks unnecessarily there is not great since the reducer actually gets us back to our initial chunk size and the splitting would lead to a highly fragmented array. My suggestion here is to expose an option for xarray to turn the rechunk off, if xarray is sure that memory won't blow up.

Thoughts @dcherian for this?




Hints:
> My suggestions would be to split up the other chunks in sliding_window_view to keep chunk sizes the same after construct, this would slow us down for cases where we indeed don't make a copy, but I think this is sufficiently rare in the general construct case.

:+1:

> Construct is also used internally for reductions, splitting up chunks unnecessarily there is not great since the reducer actually gets us back to our initial chunk size and the splitting would lead to a highly fragmented array. My suggestion here is to expose an option for xarray to turn the rechunk off, if xarray is sure that memory won't blow up.

:+1: Yes, we'd be happy to opt-in to this behaviour.

Patch:
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
index 265a50b141c..43c4ba7a252 100644
--- a/dask/array/overlap.py
+++ b/dask/array/overlap.py
@@ -1,13 +1,16 @@
 from __future__ import annotations
 
 import warnings
+from functools import reduce
 from numbers import Integral, Number
+from operator import mul
 
 import numpy as np
 from tlz import concat, get, partial
 from tlz.curried import map
 
 from dask.array import chunk
+from dask.array._shuffle import _calculate_new_chunksizes
 from dask.array.core import Array, concatenate, map_blocks, unify_chunks
 from dask.array.creation import empty_like, full_like, repeat
 from dask.array.numpy_compat import normalize_axis_tuple
@@ -805,7 +808,7 @@ def coerce_boundary(ndim, boundary):
 
 
 @derived_from(np.lib.stride_tricks)
-def sliding_window_view(x, window_shape, axis=None):
+def sliding_window_view(x, window_shape, axis=None, automatic_rechunk=True):
     window_shape = tuple(window_shape) if np.iterable(window_shape) else (window_shape,)
 
     window_shape_array = np.array(window_shape)
@@ -836,10 +839,25 @@ def sliding_window_view(x, window_shape, axis=None):
 
     # Ensure that each chunk is big enough to leave at least a size-1 chunk
     # after windowing (this is only really necessary for the last chunk).
-    safe_chunks = tuple(
+    safe_chunks = list(
         ensure_minimum_chunksize(d + 1, c) for d, c in zip(depths, x.chunks)
     )
-    x = x.rechunk(safe_chunks)
+    if automatic_rechunk:
+        safe_chunks = [
+            s if d != 0 else c for d, c, s in zip(depths, x.chunks, safe_chunks)
+        ]
+        # safe chunks is our output chunks, so add the new dimensions
+        safe_chunks.extend([(w,) for w in window_shape])
+        max_chunk = reduce(mul, map(max, x.chunks))
+        new_chunks = _calculate_new_chunksizes(
+            x.chunks,
+            safe_chunks.copy(),
+            {i for i, d in enumerate(depths) if d == 0},
+            max_chunk,
+        )
+        x = x.rechunk(tuple(new_chunks))
+    else:
+        x = x.rechunk(tuple(safe_chunks))
 
     # result.shape = x_shape_trimmed + window_shape,
     # where x_shape_trimmed is x.shape with every entry


Test patch:
diff --git a/dask/array/tests/test_overlap.py b/dask/array/tests/test_overlap.py
index beff0650e4c..2c1b2a7ee8c 100644
--- a/dask/array/tests/test_overlap.py
+++ b/dask/array/tests/test_overlap.py
@@ -831,6 +831,63 @@ def test_sliding_window_errors(window_shape, axis):
         sliding_window_view(arr, window_shape, axis)
 
 
+@pytest.mark.parametrize(
+    "output_chunks, window_shape, axis",
+    [
+        (((25, 21), (9, 8, 8, 9, 8, 8), (9, 8, 8, 9, 8, 8), (5,)), 5, 0),
+        (((25, 6), (5,) * 10, (5,) * 10, (20,)), 20, 0),
+        (
+            (
+                (11,),
+                (3, 3, 3, 3, 3, 3, 3, 2, 2) * 2,
+                (3, 3, 3, 3, 3, 3, 3, 2, 2) * 2,
+                (40,),
+            ),
+            40,
+            0,
+        ),
+        (
+            (
+                (25, 21),
+                (2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1) * 2,
+                (25, 22),
+                (5,),
+                (4,),
+            ),
+            (5, 4),
+            (0, 2),
+        ),
+        (
+            (
+                (25, 21),
+                (25, 22),
+                (2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1) * 2,
+                (5,),
+                (4,),
+            ),
+            (5, 4),
+            (0, 1),
+        ),
+    ],
+)
+def test_sliding_window_view_chunking(output_chunks, window_shape, axis):
+    arr = np.random.randn(50, 50, 50)
+    darr = da.from_array(arr, chunks=(25, 25, 25))
+    result = sliding_window_view(darr, window_shape, axis, automatic_rechunk=True)
+    assert result.chunks == output_chunks
+    expected = np.lib.stride_tricks.sliding_window_view(arr, window_shape, axis)
+    assert_eq(result, expected)
+
+
+def test_sliding_window_view_no_chunking():
+    arr = np.random.randn(50, 50, 50)
+    darr = da.from_array(arr, chunks=(25, 25, 25))
+    result = sliding_window_view(darr, 30, 0, automatic_rechunk=False)
+    assert result.chunks == ((21,), (25, 25), (25, 25), (30,))
+    expected = np.lib.stride_tricks.sliding_window_view(arr, 30, 0)
+    assert_eq(result, expected)
+
+
 def test_overlap_not_adding_empty_tasks():
     arr = da.zeros((30, 5), chunks=(10, 5))
 
@@ -875,4 +932,4 @@ def test_overlap_not_blowing_up_graph():
     )
     result = arr.rolling(dim={"year": 11}, center=True).construct("window_dim")
     dc = collections_to_dsk([result.data])
-    assert len(dc) < 1000  # previously 3000
+    assert len(dc) <= 1651  # previously 3000