Repository: dask/dask

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?



Discussion:
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     13 files  ± 0       13 suites  ±0   3h 50m 12s ⏱️ + 4m 45s
 13 259 tests + 4   12 197 ✅ + 4   1 062 💤 ±0  0 ❌ ±0 
138 342 runs  +44  119 351 ✅ +43  18 991 💤 +1  0 ❌ ±0 

Results for commit 4e11fea0. ± Comparison against base commit 0cb1cb1f.

[test-results]:data:application/gzip;base64,H4sIAG7mI2cC/03Myw6DIBCF4VcxrLtgAIHpyxiKkJB6aRBWpu9evBRd/t9kzkp8GNxCng3wR0OWHNJVfY4mhXnaWwMrVK7puLMW/90t2doNGaC68B0+G1JZHztvwlCMVnAxzvGUmKdjWnPBzq7TgLyFiue0Rqx2m977vmzncQypBBEOwDtDEQ337UsojkqgYVpICWih984DUEW+Pz6PSFIaAQAA

dcherian: +10 for on by default. Xarray can opt out intentionally where it makes sense. 

`sliding_window_view` is somewhat popular in numpy, so people will look for it in dask or dispatch to it with `__array_function__`. 

phofl: Sounds good, turned it on by default.

I'd recommend rechunking back to the original structure if you opt out of it for reductions like std. Otherwise the chunks after something like rolling().std() will be tiny
dcherian:  so... would `np.lib.stride_tricks.sliding_window_view` send the `automatic_rechunk` kwarg to dask? That's the dispatch mechanism we use today.
phofl: The sliding window view dispatcher doesn't have kwargs, so you can't pass things through unfortunately. Is there a better way of triggering this behavior to make it easier for you on the xarray side?
dcherian: I *think* this means we always have to special case dask so i don't think it matters. 

for backcompat purposes, users might appreciate using `dask.config.set` to set it?


phofl: > for backcompat purposes, users might appreciate using dask.config.set to set it?

Would this help xarray? I generally wouldn't want to make it too easy for end users to disable this behavior. Copy and Views in NumPy is a non-trivial topic that get's even more complicated with Dask 

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 3065406a9..b288d7d2d 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1111,6 +1111,7 @@ def store(
     regions: tuple[slice, ...] | Collection[tuple[slice, ...]] | None = None,
     compute: bool = True,
     return_stored: bool = False,
+    load_stored: bool | None = None,
     **kwargs,
 ):
     """Store dask arrays in array-like objects, overwrite data in target
@@ -1145,6 +1146,13 @@ def store(
         If true compute immediately; return :class:`dask.delayed.Delayed` otherwise.
     return_stored: boolean, optional
         Optionally return the stored result (default False).
+    load_stored: boolean, optional
+        Optionally return the stored result, loaded in to memory (default None).
+        If None, ``load_stored`` is True if ``return_stored`` is True and
+        ``compute`` is False. *This is an advanced option.*
+        When False, store will return the appropriate ``target`` for each chunk that is stored.
+        Directly computing this result is not what you want.
+        Instead, you can use the returned ``target`` to execute followup operations to the store.
     kwargs:
         Parameters passed to compute/persist (only used if compute=True)
 
@@ -1227,7 +1235,8 @@ def store(
         layers[targets_name] = targets_layer
         dependencies[targets_name] = set()
 
-    load_stored = return_stored and not compute
+    if load_stored is None:
+        load_stored = return_stored and not compute
 
     map_names = [
         "store-map-" + tokenize(s, t if isinstance(t, Delayed) else id(t), r)
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
index 265a50b14..43c4ba7a2 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
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index a0dfb3d4f..49bbe14d8 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -1927,7 +1927,7 @@ def _span_indexers(a):
     # we have to calculate reduce(mul, list(a.shape)[:-1]) number of quantiles
     # We create an indexer combination for each of these quantiles
 
-    shapes = reduce(mul, list(a.shape)[1:-1])
+    shapes = 1 if len(a.shape) <= 2 else reduce(mul, list(a.shape)[1:-1])
     original_shapes = shapes * a.shape[0]
     indexers = [tuple(np.repeat(np.arange(a.shape[0]), shapes))]
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index b57359bf5..2bb527cd6 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1959,13 +1959,22 @@ def test_store_compute_false():
 
     at = np.zeros(shape=(4, 4))
     bt = np.zeros(shape=(4, 4))
-
     dat, dbt = store([a, b], [at, bt], compute=False, return_stored=True)
     assert isinstance(dat, Array) and isinstance(dbt, Array)
     assert (at == 0).all() and (bt == 0).all()
     assert (dat.compute() == at).all() and (dbt.compute() == bt).all()
     assert (at == 2).all() and (bt == 3).all()
 
+    at = np.zeros(shape=(4, 4))
+    bt = np.zeros(shape=(4, 4))
+    dat, dbt = store(
+        [a, b], [at, bt], compute=False, return_stored=True, load_stored=False
+    )
+    assert isinstance(dat, Array) and isinstance(dbt, Array)
+    assert (at == 0).all() and (bt == 0).all()
+    dask.compute(dat, dbt)
+    assert (at == 2).all() and (bt == 3).all()
+
 
 def test_store_nocompute_regions():
     x = da.ones(10, chunks=1)
diff --git a/dask/array/tests/test_overlap.py b/dask/array/tests/test_overlap.py
index beff0650e..2c1b2a7ee 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
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
index 0c43c13fb..01625c0eb 100644
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -972,3 +972,9 @@ def test_nanquantile_one_dim():
     arr = np.random.randn(10)
     darr = da.from_array(arr, chunks=(2,))
     assert_eq(da.nanquantile(darr, 0.75, axis=-1), np.nanquantile(arr, 0.75, axis=-1))
+
+
+def test_nanquantile_two_dims():
+    arr = np.random.randn(10, 10)
+    darr = da.from_array(arr, chunks=(2, -1))
+    assert_eq(da.nanquantile(darr, 0.75, axis=-1), np.nanquantile(arr, 0.75, axis=-1))