Repository: dask/dask
Pull request #11465 (2024-10-29 13:59:50)

Add `load_stored` kwarg to `dask.array.store`.
Prototyping a way forward for writing to icechunk.


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 45m 13s ⏱️ -46s
 13 241 tests ±0   12 179 ✅ ±0   1 062 💤 ±0  0 ❌ ±0 
138 140 runs  ±0  119 050 ✅ +1  19 090 💤  - 1  0 ❌ ±0 

Results for commit 3ce57a03. ± Comparison against base commit 206359f4.

[test-results]:data:application/gzip;base64,H4sIALrzIGcC/03Myw6DIBCF4VcxrLtgBhGnL9PQERNSLw3CyvTd6xVd/t9Jzixa37lJPAtQj0JMycermhRs9OOwtd5oWeO+Ywlnv6bEvCKCoQs//ruirDBba323mMzgQhjDISEN+3UN5dn5GkjqC49rkpTtdr31/ZnHvvdxCaHYaWOlqoAMkzG1BKfx7ZRtiRpUaA02LFn8/rt9W/YaAQAA

rjzamora: rerun tests

Patch:
diff --git a/.github/workflows/additional.yml b/.github/workflows/additional.yml
index 78bd40f9e..8e4d0467a 100644
--- a/.github/workflows/additional.yml
+++ b/.github/workflows/additional.yml
@@ -16,7 +16,7 @@ jobs:
         uses: actions/checkout@v4.1.5
 
       - name: Setup Conda Environment
-        uses: conda-incubator/setup-miniconda@v3.0.4
+        uses: conda-incubator/setup-miniconda@v3.1.0
         with:
           miniforge-version: latest
           use-mamba: true
@@ -46,7 +46,7 @@ jobs:
         uses: actions/checkout@v4.1.5
 
       - name: Setup Conda
-        uses: conda-incubator/setup-miniconda@v3.0.4
+        uses: conda-incubator/setup-miniconda@v3.1.0
         with:
           miniforge-version: latest
           use-mamba: true
diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml
index ca93fa2cc..fe2b4b8e7 100644
--- a/.github/workflows/conda.yml
+++ b/.github/workflows/conda.yml
@@ -32,7 +32,7 @@ jobs:
         with:
           fetch-depth: 0
       - name: Set up Python
-        uses: conda-incubator/setup-miniconda@v3.0.4
+        uses: conda-incubator/setup-miniconda@v3.1.0
         with:
           miniforge-version: latest
           use-mamba: true
diff --git a/.github/workflows/upstream.yml b/.github/workflows/upstream.yml
index 07d7ebbc6..2b835b517 100644
--- a/.github/workflows/upstream.yml
+++ b/.github/workflows/upstream.yml
@@ -50,7 +50,7 @@ jobs:
         uses: actions/checkout@v4.1.5
 
       - name: Setup Conda Environment
-        uses: conda-incubator/setup-miniconda@v3.0.4
+        uses: conda-incubator/setup-miniconda@v3.1.0
         with:
           miniforge-version: latest
           use-mamba: true
diff --git a/dask/array/core.py b/dask/array/core.py
index fdf65bd24..b288d7d2d 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -562,7 +562,9 @@ def map_blocks(
         Dimensions lost by the function.
     new_axis : number or iterable, optional
         New dimensions created by the function. Note that these are applied
-        after ``drop_axis`` (if present).
+        after ``drop_axis`` (if present). The size of each chunk along this
+        dimension will be set to 1. Please specify ``chunks`` if the individual
+        chunks have a different size.
     enforce_ndim : bool, default False
         Whether to enforce at runtime that the dimensionality of the array
         produced by ``func`` actually matches that of the array returned by
@@ -1109,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
@@ -1143,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)
 
@@ -1225,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/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/.github/workflows/test-report.yaml b/.github/workflows/test-report.yaml
index 615363e6d..c84889062 100644
--- a/.github/workflows/test-report.yaml
+++ b/.github/workflows/test-report.yaml
@@ -26,7 +26,7 @@ jobs:
           repository: dask/distributed
 
       - name: Setup Conda Environment
-        uses: conda-incubator/setup-miniconda@v3.0.4
+        uses: conda-incubator/setup-miniconda@v3.1.0
         with:
           miniforge-version: latest
           condarc-file: continuous_integration/condarc
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 444a35fbb..fe76623df 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -69,7 +69,7 @@ jobs:
           java-version: "11"
 
       - name: Setup Conda Environment
-        uses: conda-incubator/setup-miniconda@v3.0.4
+        uses: conda-incubator/setup-miniconda@v3.1.0
         with:
           miniforge-version: latest
           use-mamba: true
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_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))