Repository: dask/dask
Pull request #11548 (2024-11-20 17:52:57)

Migrate slicing code to task spec
- [ ] 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  +     10       13 suites  +10   3h 16m 21s ⏱️ + 3h 8m 52s
 13 255 tests +  7 806   12 192 ✅ +  7 541   1 062 💤 +   264   1 ❌ + 1 
138 315 runs  +130 443  119 306 ✅ +112 691  18 998 💤 +17 741  11 ❌ +11 

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

Results for commit 51477d34. ± Comparison against base commit b63ca061.

[test-results]:data:application/gzip;base64,H4sIAPUpPmcC/03NSw7DIAxF0a1UjDvAfGLoZqpAHAk1n4rAqOreC0pIOrzH0vOHjWGijT1uIO83tuWQrhpy7FNYl9qABgqVa9rvQuvWzy17X1GAFRe+wrsi7y4b+zBVO4FiXGMRXiTmZZ82EvTR5zRYybsTj2ljrWnWpqHB/7Rf5zmkEkyDQhykciSc8hoVEFfeyfIAvXBk0OBIJNj3B+/vQzobAQAA


Patch:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 7ae19568d..2c926d317 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -10,6 +10,7 @@ from typing import Literal
 import numpy as np
 
 from dask import config
+from dask._task_spec import DataNode, List, Task, TaskRef
 from dask.array.chunk import getitem
 from dask.array.core import Array, unknown_chunk_message
 from dask.array.dispatch import concatenate_lookup, take_lookup
@@ -226,16 +227,17 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
     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):
-        old_blocks[old_index] = (in_name,) + old_index
+    old_blocks = {
+        old_index: (in_name,) + old_index
+        for old_index in np.ndindex(tuple([len(c) for c in chunks]))
+    }
 
     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)
+        merges[sorter_key] = DataNode(sorter_key, (1, sorter))
 
         sorted_array = new_chunk_taker[sorter]
         source_chunk_nr, taker_boundary = np.unique(
@@ -271,22 +273,30 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
 
                     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))
+                    intermediates[taker_key] = DataNode(
+                        taker_key, (1, tuple(this_slice))
+                    )
                     taker_cache[c] = taker_key
 
-                intermediates[name] = _getitem, old_blocks[chunk_key], taker_key
+                intermediates[name] = Task(
+                    name, _getitem, TaskRef(old_blocks[chunk_key]), TaskRef(taker_key)
+                )
                 merge_keys.append(name)
 
             merge_suffix = convert_key(chunk_tuple, new_chunk_idx, axis)
+            out_name_merge = (out_name,) + merge_suffix
             if len(merge_keys) > 1:
-                merges[(out_name,) + merge_suffix] = (
+                merges[out_name_merge] = Task(
+                    out_name_merge,
                     concatenate_arrays,
-                    merge_keys,
-                    sorter_key,
+                    List(*(TaskRef(m) for m in merge_keys)),
+                    TaskRef(sorter_key),
                     axis,
                 )
             elif len(merge_keys) == 1:
-                merges[(out_name,) + merge_suffix] = intermediates.pop(merge_keys[0])
+                t = intermediates.pop(merge_keys[0])
+                t.key = out_name_merge
+                merges[out_name_merge] = t
             else:
                 raise NotImplementedError
 
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index ec12479cf..d3bd3a9d9 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -12,6 +12,7 @@ import numpy as np
 from tlz import concat, memoize, merge, pluck
 
 from dask import core
+from dask._task_spec import Alias, Task, TaskRef
 from dask.array.chunk import getitem
 from dask.base import is_dask_collection, tokenize
 from dask.highlevelgraph import HighLevelGraph
@@ -141,10 +142,8 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
     >>> dsk, blockdims = slice_array('y', 'x', [(20, 20, 20, 20, 20)],
     ...                              (slice(10, 35),), 8)
     >>> pprint(dsk)  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
-    {('y', 0): (<function getitem at ...>,
-                ('x', 0),
-                (slice(10, 20, 1),)),
-     ('y', 1): (<function getitem at ...>, ('x', 1), (slice(0, 15, 1),))}
+    {('y', 0): <Task ('y', 0) getitem(Alias(('x', 0)->('x', 0)), (slice(10, 20, 1),))>,
+     ('y', 1): <Task ('y', 1) getitem(Alias(('x', 1)->('x', 1)), (slice(0, 15, 1),))>}
     >>> blockdims
     ((10, 15),)
 
@@ -164,7 +163,9 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
         isinstance(index, slice) and index == slice(None, None, None) for index in index
     ):
         suffixes = product(*[range(len(bd)) for bd in blockdims])
-        dsk = {(out_name,) + s: (in_name,) + s for s in suffixes}
+        dsk = {
+            (out_name,) + s: Alias((out_name,) + s, (in_name,) + s) for s in suffixes
+        }
         return dsk, blockdims
 
     # Add in missing colons at the end as needed.  x[5] -> x[5, :, :]
@@ -204,19 +205,17 @@ def slice_with_newaxes(out_name, in_name, blockdims, index):
         expand_orig = expander(where_none_orig)
 
         # Insert ",0" into the key:  ('x', 2, 3) -> ('x', 0, 2, 0, 3)
-        dsk2 = {
-            (out_name,) + expand(k[1:], 0): (v[:2] + (expand_orig(v[2], None),))
-            for k, v in dsk.items()
-            if k[0] == out_name
-        }
-
-        # Add back intermediate parts of the dask that weren't the output
-        dsk3 = merge(dsk2, {k: v for k, v in dsk.items() if k[0] != out_name})
+        dsk2 = {}
+        for k, v in dsk.items():
+            if k[0] == out_name:
+                k2 = (out_name,) + expand(k[1:], 0)
+                dsk2[k2] = Task(k2, v.func, v.args[0], expand_orig(v.args[1], None))
+            else:
+                dsk2[k] = v
 
         # Insert (1,) into blockdims:  ((2, 2), (3, 3)) -> ((2, 2), (1,), (3, 3))
         blockdims3 = expand(blockdims2, (1,))
-
-        return dsk3, blockdims3
+        return dsk2, blockdims3
 
     else:
         return dsk, blockdims2
@@ -340,10 +339,10 @@ def slice_slices_and_integers(
 
     dsk_out = {
         out_name: (
-            (getitem, in_name, slices)
+            Task(out_name, getitem, TaskRef(in_name), slices)
             if not allow_getitem_optimization
             or not all(sl == slice(None, None, None) for sl in slices)
-            else in_name
+            else Alias(out_name, in_name)
         )
         for out_name, in_name, slices in zip(out_names, in_names, all_slices)
     }
@@ -587,7 +586,10 @@ def take(outname, inname, chunks, index, axis=0):
             # TODO: This should be a real no-op, but the call stack is
             # too deep to do this efficiently for now
             chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(chunks))))
-            graph = {(outname,) + c: (inname,) + c for c in chunk_tuples}
+            graph = {
+                (outname,) + c: Alias((outname,) + c, (inname,) + c)
+                for c in chunk_tuples
+            }
             return tuple(chunks), graph
 
         average_chunk_size = int(sum(chunks[axis]) / len(chunks[axis]))
@@ -610,7 +612,9 @@ def take(outname, inname, chunks, index, axis=0):
         slices = tuple(slices)
         chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(chunks))))
         dsk = {
-            (outname,) + ct: (getitem, (inname,) + ct, slices) for ct in chunk_tuples
+            (outname,)
+            + ct: Task((outname,) + ct, getitem, TaskRef((inname,) + ct), slices)
+            for ct in chunk_tuples
         }
         return chunks, dsk
     else:
diff --git a/dask/layers.py b/dask/layers.py
index c0b1384f3..5fe1bb482 100644
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -110,7 +110,7 @@ class ArrayValuesDep(ArrayBlockwiseDep):
 
 @normalize_token.register(ArraySliceDep)
 def normalize_array_slice_dep(dep):
-    return "ArraySliceDep", dep.chunks, dep.starts
+    return "ArraySliceDep", dep.chunks
 
 
 @normalize_token.register(ArrayBlockIdDep)


Test patch:
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 0b4aebbc4..fda57f41f 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][1].dtype == np.uint8
+        v.value[1][1].dtype == np.uint8
         for k, v in dict(result.dask).items()
         if "shuffle-taker" in k
     )
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index 2d1ab3006..509c2e3aa 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -5,6 +5,8 @@ import warnings
 
 import pytest
 
+from dask._task_spec import Alias, Task, TaskRef
+
 np = pytest.importorskip("numpy")
 
 import dask
@@ -177,10 +179,10 @@ def test_slice_singleton_value_on_boundary():
 def test_slice_array_1d():
     # x[24::2]
     expected = {
-        ("y", 0): (getitem, ("x", 0), (slice(24, 25, 2),)),
-        ("y", 1): (getitem, ("x", 1), (slice(1, 25, 2),)),
-        ("y", 2): (getitem, ("x", 2), (slice(0, 25, 2),)),
-        ("y", 3): (getitem, ("x", 3), (slice(1, 25, 2),)),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 0)), (slice(24, 25, 2),)),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 1)), (slice(1, 25, 2),)),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
+        ("y", 3): Task(("y", 3), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
     result, chunks = slice_array("y", "x", [[25] * 4], [slice(24, None, 2)], 8)
 
@@ -188,9 +190,9 @@ def test_slice_array_1d():
 
     # x[26::2]
     expected = {
-        ("y", 0): (getitem, ("x", 1), (slice(1, 25, 2),)),
-        ("y", 1): (getitem, ("x", 2), (slice(0, 25, 2),)),
-        ("y", 2): (getitem, ("x", 3), (slice(1, 25, 2),)),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 1)), (slice(1, 25, 2),)),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
 
     result, chunks = slice_array("y", "x", [[25] * 4], [slice(26, None, 2)], 8)
@@ -198,10 +200,10 @@ def test_slice_array_1d():
 
     # x[24::2]
     expected = {
-        ("y", 0): (getitem, ("x", 0), (slice(24, 25, 2),)),
-        ("y", 1): (getitem, ("x", 1), (slice(1, 25, 2),)),
-        ("y", 2): (getitem, ("x", 2), (slice(0, 25, 2),)),
-        ("y", 3): (getitem, ("x", 3), (slice(1, 25, 2),)),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 0)), (slice(24, 25, 2),)),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 1)), (slice(1, 25, 2),)),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
+        ("y", 3): Task(("y", 3), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
     result, chunks = slice_array("y", "x", [(25,) * 4], (slice(24, None, 2),), 8)
 
@@ -209,9 +211,9 @@ def test_slice_array_1d():
 
     # x[26::2]
     expected = {
-        ("y", 0): (getitem, ("x", 1), (slice(1, 25, 2),)),
-        ("y", 1): (getitem, ("x", 2), (slice(0, 25, 2),)),
-        ("y", 2): (getitem, ("x", 3), (slice(1, 25, 2),)),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 1)), (slice(1, 25, 2),)),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
 
     result, chunks = slice_array("y", "x", [(25,) * 4], (slice(26, None, 2),), 8)
@@ -221,16 +223,17 @@ def test_slice_array_1d():
 def test_slice_array_2d():
     # 2d slices: x[13::2,10::1]
     expected = {
-        ("y", 0, 0): (getitem, ("x", 0, 0), (slice(13, 20, 2), slice(10, 20, 1))),
-        ("y", 0, 1): (
+        ("y", 0, 0): Task(
+            ("y", 0, 0),
             getitem,
-            ("x", 0, 1),
-            (slice(13, 20, 2), slice(None, None, None)),
+            TaskRef(("x", 0, 0)),
+            (slice(13, 20, 2), slice(10, 20, 1)),
         ),
-        ("y", 0, 2): (
-            getitem,
-            ("x", 0, 2),
-            (slice(13, 20, 2), slice(None, None, None)),
+        ("y", 0, 1): Task(
+            ("y", 0, 1), getitem, TaskRef(("x", 0, 1)), (slice(13, 20, 2), slice(None))
+        ),
+        ("y", 0, 2): Task(
+            ("y", 0, 2), getitem, TaskRef(("x", 0, 2)), (slice(13, 20, 2), slice(None))
         ),
     }
 
@@ -246,9 +249,9 @@ def test_slice_array_2d():
 
     # 2d slices with one dimension: x[5,10::1]
     expected = {
-        ("y", 0): (getitem, ("x", 0, 0), (5, slice(10, 20, 1))),
-        ("y", 1): (getitem, ("x", 0, 1), (5, slice(None, None, None))),
-        ("y", 2): (getitem, ("x", 0, 2), (5, slice(None, None, None))),
+        ("y", 0): Task(("y", 0), getitem, TaskRef(("x", 0, 0)), (5, slice(10, 20, 1))),
+        ("y", 1): Task(("y", 1), getitem, TaskRef(("x", 0, 1)), (5, slice(None))),
+        ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 0, 2)), (5, slice(None))),
     }
 
     result, chunks = slice_array(
@@ -260,12 +263,16 @@ def test_slice_array_2d():
 
 def test_slice_optimizations():
     # bar[:]
-    expected = {("foo", 0): ("bar", 0)}
+    expected = {("foo", 0): Alias(("foo", 0), ("bar", 0))}
     result, chunks = slice_array("foo", "bar", [[100]], (slice(None, None, None),), 8)
     assert expected == result
 
     # bar[:,:,:]
-    expected = {("foo", 0): ("bar", 0), ("foo", 1): ("bar", 1), ("foo", 2): ("bar", 2)}
+    expected = {
+        ("foo", 0): Alias(("foo", 0), ("bar", 0)),
+        ("foo", 1): Alias(("foo", 1), ("bar", 1)),
+        ("foo", 2): Alias(("foo", 2), ("bar", 2)),
+    }
     result, chunks = slice_array(
         "foo",
         "bar",
@@ -281,7 +288,11 @@ def test_slicing_with_singleton_indices():
         "y", "x", ([5, 5], [5, 5]), (slice(0, 5), 8), itemsize=8
     )
 
-    expected = {("y", 0): (getitem, ("x", 0, 1), (slice(None, None, None), 3))}
+    expected = {
+        ("y", 0): Task(
+            ("y", 0), getitem, TaskRef(("x", 0, 1)), (slice(None, None, None), 3)
+        )
+    }
 
     assert expected == result
 
@@ -296,14 +307,16 @@ def test_slicing_with_newaxis():
     )
 
     expected = {
-        ("y", 0, 0, 0): (
+        ("y", 0, 0, 0): Task(
+            ("y", 0, 0, 0),
             getitem,
-            ("x", 0, 0),
+            TaskRef(("x", 0, 0)),
             (slice(0, 3, 1), None, slice(None, None, None)),
         ),
-        ("y", 0, 0, 1): (
+        ("y", 0, 0, 1): Task(
+            ("y", 0, 0, 1),
             getitem,
-            ("x", 0, 1),
+            TaskRef(("x", 0, 1)),
             (slice(0, 3, 1), None, slice(None, None, None)),
         ),
     }
@@ -450,15 +463,19 @@ def test_slicing_exhaustively():
 def test_slicing_with_negative_step_flops_keys():
     x = da.arange(10, chunks=5)
     y = x[:1:-1]
-    assert (x.name, 1) in y.dask[(y.name, 0)]
-    assert (x.name, 0) in y.dask[(y.name, 1)]
+    assert (x.name, 1) in y.dask[(y.name, 0)].dependencies
+    assert (x.name, 0) in y.dask[(y.name, 1)].dependencies
 
     assert_eq(y, np.arange(10)[:1:-1])
 
     assert y.chunks == ((5, 3),)
 
-    assert y.dask[(y.name, 0)] == (getitem, (x.name, 1), (slice(-1, -6, -1),))
-    assert y.dask[(y.name, 1)] == (getitem, (x.name, 0), (slice(-1, -4, -1),))
+    assert y.dask[(y.name, 0)] == Task(
+        (y.name, 0), getitem, TaskRef((x.name, 1)), (slice(-1, -6, -1),)
+    )
+    assert y.dask[(y.name, 1)] == Task(
+        (y.name, 1), getitem, TaskRef((x.name, 0)), (slice(-1, -4, -1),)
+    )
 
 
 def test_empty_slice():
@@ -559,8 +576,18 @@ def test_uneven_blockdims():
     index = (slice(240, 270), slice(None))
     dsk_out, bd_out = slice_array("in", "out", blockdims, index, itemsize=8)
     sol = {
-        ("in", 0, 0): (getitem, ("out", 7, 0), (slice(28, 31, 1), slice(None))),
-        ("in", 1, 0): (getitem, ("out", 8, 0), (slice(0, 27, 1), slice(None))),
+        ("in", 0, 0): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 7, 0)),
+            (slice(28, 31, 1), slice(None)),
+        ),
+        ("in", 1, 0): Task(
+            ("in", 1, 0),
+            getitem,
+            TaskRef(("out", 8, 0)),
+            (slice(0, 27, 1), slice(None)),
+        ),
     }
     assert dsk_out == sol
     assert bd_out == ((3, 27), (100,))
@@ -569,12 +596,42 @@ def test_uneven_blockdims():
     index = (slice(240, 270), slice(180, 230))
     dsk_out, bd_out = slice_array("in", "out", blockdims, index, itemsize=8)
     sol = {
-        ("in", 0, 0): (getitem, ("out", 7, 5), (slice(28, 31, 1), slice(29, 30, 1))),
-        ("in", 0, 1): (getitem, ("out", 7, 6), (slice(28, 31, 1), slice(None))),
-        ("in", 0, 2): (getitem, ("out", 7, 7), (slice(28, 31, 1), slice(0, 18, 1))),
-        ("in", 1, 0): (getitem, ("out", 8, 5), (slice(0, 27, 1), slice(29, 30, 1))),
-        ("in", 1, 1): (getitem, ("out", 8, 6), (slice(0, 27, 1), slice(None))),
-        ("in", 1, 2): (getitem, ("out", 8, 7), (slice(0, 27, 1), slice(0, 18, 1))),
+        ("in", 0, 0): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 7, 5)),
+            (slice(28, 31, 1), slice(29, 30, 1)),
+        ),
+        ("in", 0, 1): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 7, 6)),
+            (slice(28, 31, 1), slice(None)),
+        ),
+        ("in", 0, 2): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 7, 7)),
+            (slice(28, 31, 1), slice(0, 18, 1)),
+        ),
+        ("in", 1, 0): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 8, 5)),
+            (slice(0, 27, 1), slice(29, 30, 1)),
+        ),
+        ("in", 1, 1): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 8, 6)),
+            (slice(0, 27, 1), slice(None)),
+        ),
+        ("in", 1, 2): Task(
+            ("in", 0, 0),
+            getitem,
+            TaskRef(("out", 8, 7)),
+            (slice(0, 27, 1), slice(0, 18, 1)),
+        ),
     }
     assert dsk_out == sol
     assert bd_out == ((3, 27), (1, 31, 18))
@@ -1029,7 +1086,7 @@ def test_take_sorted_indexer():
     assert {
         **dict(arr.dask),
         **{
-            k: k2
+            k: Alias(k, k2)
             for k, k2 in zip(
                 [k for k in dict(result.dask) if "getitem" in k[0]],
                 dict(arr.dask).keys(),
@@ -1044,5 +1101,5 @@ def test_all_none_slices_just_mappings():
     dsk = dict(result.dask)
     assert len([k for k in dsk if "getitem" in k[0]]) == 12
     # check that we are just mapping the keys
-    assert all(v in dsk for k, v in dsk.items() if "getitem" in k[0])
+    assert all(isinstance(v, Alias) for k, v in dsk.items() if "getitem" in k[0])
     assert_eq(result, np.ones((6, 10)))