Repository: dask/dask
Pull request #11362 (2024-09-02 14:50:29)

Write indexing arrays into dask graph to reduce size for multiple xarray variables
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


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 4m 48s :stopwatch: -30s
 13 170 tests +1   12 113 :white_check_mark: ±0   1 056 :zzz: ±0  1 :x: +1 
137 254 runs  +8  118 221 :white_check_mark: ±0  19 025 :zzz: ±0  8 :x: +8 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/29572388110).

Results for commit 30329b14. ± Comparison against base commit 80b0737c.

[test-results]:data:application/gzip;base64,H4sIAHHY1WYC/03MTQ6DIBCG4asY1l0wIwj2Mo2DmJD60yCsTO9eqEi7fJ/JfAeb3Gx3dm+gvTVsjy78aox+CG5bcwPXOlG6hvMOil/92KMxGRG+fwWf7pWRy67aNLg5WwXr/eaT5C0f13NaoRSl6zRoRKhYpnuO8rIyra/+XzbbsriQgrW8xZ5AEGhhUXLixnZCWYPUk6JpJGNQILH3B6BkAwcaAQAA

phofl: AB tests look good https://github.com/coiled/benchmarks/actions/runs/10683023390/artifacts/1886066719, I don't trust the speedup though, that seems a bit odd

Patch:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 25320100d..ad11f299a 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -211,6 +211,8 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
     split_name = f"shuffle-split-{token}"
     slices = [slice(None)] * len(chunks)
     split_name_suffixes = count()
+    sorter_name = "shuffle-sorter-"
+    taker_name = "shuffle-taker-"
 
     old_blocks = np.empty([len(c) for c in chunks], dtype="O")
     for old_index in np.ndindex(old_blocks.shape):
@@ -219,6 +221,10 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
     for new_chunk_idx, new_chunk_taker in enumerate(new_chunks):
         new_chunk_taker = np.array(new_chunk_taker)
         sorter = np.argsort(new_chunk_taker).astype(dtype)
+        sorter_key = sorter_name + tokenize(sorter)
+        # low level fusion can't deal with arrays on first position
+        merges[sorter_key] = (1, sorter)
+
         sorted_array = new_chunk_taker[sorter]
         source_chunk_nr, taker_boundary = np.unique(
             np.searchsorted(chunk_boundaries, sorted_array, side="right"),
@@ -242,7 +248,7 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
                 # Cache the takers to allow de-duplication when serializing
                 # Ugly!
                 if c in taker_cache:
-                    this_slice[axis] = taker_cache[c]
+                    taker_key = taker_cache[c]
                 else:
                     this_slice[axis] = (
                         sorted_array[b_start:b_end]
@@ -250,9 +256,13 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
                     ).astype(dtype)
                     if len(source_chunk_nr) == 1:
                         this_slice[axis] = this_slice[axis][np.argsort(sorter)]
-                    taker_cache[c] = this_slice[axis]
 
-                intermediates[name] = getitem, old_blocks[chunk_key], tuple(this_slice)
+                    taker_key = taker_name + tokenize(this_slice)
+                    # low level fusion can't deal with arrays on first position
+                    intermediates[taker_key] = (1, tuple(this_slice))
+                    taker_cache[c] = taker_key
+
+                intermediates[name] = _getitem, old_blocks[chunk_key], taker_key
                 merge_keys.append(name)
 
             merge_suffix = convert_key(chunk_tuple, new_chunk_idx, axis)
@@ -260,7 +270,7 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
                 merges[(out_name,) + merge_suffix] = (
                     concatenate_arrays,
                     merge_keys,
-                    sorter,
+                    sorter_key,
                     axis,
                 )
             elif len(merge_keys) == 1:
@@ -279,9 +289,13 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
     return tuple(output_chunks), layer
 
 
+def _getitem(obj, index):
+    return getitem(obj, index[1])
+
+
 def concatenate_arrays(arrs, sorter, axis):
     concatenate = concatenate_lookup.dispatch(type(arrs[0]))
-    return np.take(concatenate(arrs, axis=axis), np.argsort(sorter), axis=axis)
+    return np.take(concatenate(arrs, axis=axis), np.argsort(sorter[1]), axis=axis)
 
 
 def convert_key(key, chunk, axis):
diff --git a/dask/array/utils.py b/dask/array/utils.py
index b515f8b84..1c7adc325 100644
--- a/dask/array/utils.py
+++ b/dask/array/utils.py
@@ -213,13 +213,18 @@ def _check_dsk(dsk):
     non_one = {k: v for k, v in freqs.items() if v != 1}
     key_collisions = set()
     # Allow redundant keys if the values are equivalent
+    collisions = set()
     for k in non_one.keys():
         for layer in dsk.layers.values():
             try:
                 key_collisions.add(tokenize(layer[k]))
             except KeyError:
                 pass
-    assert len(key_collisions) < 2, non_one
+        if len(key_collisions) >= 2:
+            collisions.add(k)
+        key_collisions.clear()
+
+    assert len(collisions) == 0, {k: non_one[k] for k in collisions}
 
 
 def assert_eq_shape(a, b, check_ndim=True, check_nan=True):


Test patch:
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 70339768b..0b4aebbc4 100644
--- a/dask/array/tests/test_shuffle.py
+++ b/dask/array/tests/test_shuffle.py
@@ -148,7 +148,7 @@ def test_dtype_taker(arr, darr):
     expected = arr[:, [1, 5, 6, 0, 3, 4, 2, 7]]
     assert_eq(result, expected)
     assert all(
-        v[1].dtype == np.uint8
+        v[1][1].dtype == np.uint8
         for k, v in dict(result.dask).items()
-        if "array-taker" in k
+        if "shuffle-taker" in k
     )
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index f503c5fee..a5ea27370 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -5,8 +5,6 @@ import warnings
 
 import pytest
 
-from dask.array._shuffle import concatenate_arrays
-
 np = pytest.importorskip("numpy")
 
 import dask
@@ -316,109 +314,21 @@ def test_slicing_with_newaxis():
 
 def test_take():
     chunks, dsk = take("y-y", "x", [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
-    expected = {
-        ("y-y", 0): (
-            concatenate_arrays,
-            [
-                ("shuffle-split-y", 0),
-                ("shuffle-split-y", 1),
-            ],
-            np.array([1, 3, 0, 2]),
-            0,
-        ),
-        ("shuffle-split-y", 0): (
-            getitem,
-            ("x", 0),
-            (np.array([1, 3, 5]),),
-        ),
-        ("shuffle-split-y", 1): (
-            getitem,
-            ("x", 2),
-            (np.array([7]),),
-        ),
-    }
-    np.testing.assert_equal(sorted(dsk.items()), sorted(expected.items()))
+    assert len(dsk) == 6
     assert chunks == ((4,),)
 
     chunks, dsk = take("y-y", "x", [(20, 20, 20, 20), (20, 20)], [5, 1, 47, 3], axis=0)
-    expected = {
-        ("y-y", 0, 0): (
-            concatenate_arrays,
-            [
-                ("shuffle-split-y", 0),
-                ("shuffle-split-y", 1),
-            ],
-            np.array([1, 3, 0, 2]),
-            0,
-        ),
-        ("y-y", 0, 1): (
-            concatenate_arrays,
-            [
-                ("shuffle-split-y", 2),
-                ("shuffle-split-y", 3),
-            ],
-            np.array([1, 3, 0, 2]),
-            0,
-        ),
-        ("shuffle-split-y", 0): (
-            getitem,
-            ("x", 0, 0),
-            (np.array([1, 3, 5]), slice(None)),
-        ),
-        ("shuffle-split-y", 1): (
-            getitem,
-            ("x", 2, 0),
-            (np.array([7]), slice(None)),
-        ),
-        ("shuffle-split-y", 2): (
-            getitem,
-            ("x", 0, 1),
-            (np.array([1, 3, 5]), slice(None)),
-        ),
-        ("shuffle-split-y", 3): (
-            getitem,
-            ("x", 2, 1),
-            (np.array([7]), slice(None)),
-        ),
-    }
-    np.testing.assert_equal(sorted(dsk.items()), sorted(expected.items()))
+    assert len(dsk) == 9
     assert chunks == ((4,), (20, 20))
 
 
 def test_take_sorted():
     chunks, dsk = take("y-y", "x", [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
-    expected = {
-        ("y-y", 0): (
-            concatenate_arrays,
-            [
-                ("shuffle-split-y", 0),
-                ("shuffle-split-y", 1),
-            ],
-            np.array([0, 1, 2, 3]),
-            0,
-        ),
-        ("shuffle-split-y", 0): (
-            getitem,
-            ("x", 0),
-            (np.array([1, 3, 5]),),
-        ),
-        ("shuffle-split-y", 1): (
-            getitem,
-            ("x", 2),
-            (np.array([7]),),
-        ),
-    }
-    np.testing.assert_equal(dsk, expected)
+    assert len(dsk) == 6
     assert chunks == ((4,),)
 
     chunks, dsk = take("y", "x", [(20, 20, 20, 20)], np.arange(0, 80), axis=0)
-    expected = {
-        ("y", 0): ("x", 0),
-        ("y", 1): ("x", 1),
-        ("y", 2): ("x", 2),
-        ("y", 3): ("x", 3),
-    }
-    np.testing.assert_equal(dsk, expected)
+    assert len(dsk) == 4
     assert chunks == ((20, 20, 20, 20),)
 
 
@@ -964,23 +874,23 @@ def test_take_avoids_large_chunks():
         index = np.array([0, 1] + [2] * 101 + [3])
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
         index = np.array([0] * 101 + [1, 2, 3])
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
         index = np.array([0, 1, 2] + [3] * 101)
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
         chunks = ((500,), (1, 1, 1, 1), (500,))
         index = np.array([0, 1, 2] + [3] * 101)
         chunks2, dsk = take("a", "b", chunks, index, axis=1)
         assert chunks2 == ((500,), (1,) * 104, (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
 
 def test_take_uses_config():
@@ -990,7 +900,7 @@ def test_take_uses_config():
         with config.set({"array.chunk-size": "10GB"}):
             chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
 
 def test_pathological_unsorted_slicing():
@@ -1048,6 +958,13 @@ def test_shuffle_slice(size, chunks):
     b = shuffle_slice(x, index)
     assert_eq(a, b)
 
+    index = np.arange(1, len(x)).tolist()
+    index.append(0)
+    index = np.array(index)
+    a = x[index]
+    b = shuffle_slice(x, index)
+    assert_eq(a, b)
+
 
 def test_unknown_chunks_length_one():
     a = np.arange(256, dtype=int)
diff --git a/dask/array/tests/test_xarray.py b/dask/array/tests/test_xarray.py
index 4bb0739d9..9740996d5 100644
--- a/dask/array/tests/test_xarray.py
+++ b/dask/array/tests/test_xarray.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import numpy as np
 import pytest
 
 import dask.array as da
@@ -54,3 +55,30 @@ def test_polyfit_reshaping():
     arr = xr.DataArray(da.ones((10, 20, 30), chunks=(1, 5, 30)), dims=["z", "y", "x"])
     result = arr.polyfit("x", 1)
     assert result.polyfit_coefficients.chunks == ((2,), (1,) * 10, (5,) * 4)
+
+
+def test_positional_indexer_multiple_variables():
+    n = 200
+    ds = xr.Dataset(
+        data_vars=dict(
+            a=(
+                ["x", "y", "time"],
+                da.random.randint(1, 100, (10, 10, n), chunks=(-1, -1, n // 2)),
+            ),
+            b=(
+                ["x", "y", "time"],
+                da.random.randint(1, 100, (10, 10, n), chunks=(-1, -1, n // 2)),
+            ),
+        ),
+        coords=dict(
+            x=list(range(10)),
+            y=list(range(10)),
+            time=np.arange(n),
+        ),
+    )
+    indexer = np.arange(n)
+    np.random.shuffle(indexer)
+    result = ds.isel(time=indexer)
+    graph = result.__dask_graph__()
+    assert len({k for k in graph if "shuffle-taker" in k}) == 4
+    assert len({k for k in graph if "shuffle-sorter" in k}) == 2
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index e05252c62..ea9326b52 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -5523,7 +5523,6 @@ def test_broadcast():
     assert_eq(ddf - (ddf.sum() + 1), df - (df.sum() + 1))
 
 
-@pytest.mark.xfail(DASK_EXPR_ENABLED, reason="array computation doesn't work yet")
 def test_scalar_with_array():
     df = pd.DataFrame({"x": [1, 2, 3, 4, 5]})
     ddf = dd.from_pandas(df, npartitions=2)
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 008795b2f..e86a711c9 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -1254,13 +1254,15 @@ def test_anom_mean():
     # `mean_chunk` which is the primary reducer in this graph. Therefore we want
     # to run those as quickly as possible.
     # This is difficult to assert on but the pressure is an ok-ish proxy
-    assert max(pressure) <= 177
+    assert max(pressure) <= 178
     from collections import defaultdict
 
     count_dependents = defaultdict(set)
     for k in dict(graph).keys():
         count_dependents[len(dependents[k])].add(k)
-    n_splits = max(count_dependents)
+
+    # array-taker has the most dependents, but it's not what we want to look at
+    n_splits = sorted(count_dependents)[-2]
     # There is a transpose/stack group that is splitting into many tasks
     # see https://github.com/dask/dask/pull/10660#discussion_r1420571664
     # the name depends on the version of xarray
@@ -1285,7 +1287,7 @@ def test_anom_mean():
     max_age_transpose = max(ages_tranpose.values())
     assert max_age_transpose < 150
     assert avg_age_transpose < 100
-    assert sum(pressure) / len(pressure) < 100
+    assert sum(pressure) / len(pressure) < 101
 
 
 def test_anom_mean_raw(abcde):