Repository: dask/dask
Pull request #11946 (2025-05-14 13:42:20)

Optimize slicing graph generation
This shaves off another 20-30% of the slicing case presented in https://github.com/dask/dask/issues/11735#issuecomment-2876708328

On this PR I get something like 900-1000ms
Main is somewhere between 1200-1300ms
Pre-Task-class: ~750ms

So, we're still about 70% slower than the 2024.X versions. I ran some more realistic end-to-end benchmarks that includes optimization, dependency calculation and materialization and with this we're either equally slow or faster than we were on the 2024.11.2 version.

```
Dask version: 2024.11.2
Empty indexer (all)                     :    2086.02 ms
Every second element (::2)              :    2503.59 ms
All but last element (:-1)              :    4553.73 ms
Multi dim one slice (0, slice(None))    :    5283.48 ms
Multi dim (0, slice(None, None, 2))     :    2486.43 ms
Multi dim (0, slice(None, -1))          :    4965.22 ms
Multi dim (slice(None, -1), slice(None, 10)):   11988.26 ms
(dask-distributed)
~/workspace/dask  f0e08da76 ✗                                                                                                                                                                                                      181d22h ◒

# This PR
▶ python profile_slicing.py
Dask version: 2025.4.1+16.g0a0f974b2
Empty indexer (all)                     :    2196.06 ms
Every second element (::2)              :    2177.34 ms
All but last element (:-1)              :    4955.38 ms
Multi dim one slice (0, slice(None))    :    4464.50 ms
Multi dim (0, slice(None, None, 2))     :    2201.39 ms
Multi dim (0, slice(None, -1))          :    4561.45 ms
Multi dim (slice(None, -1), slice(None, 10)):   10419.56 ms

```

```python
from __future__ import annotations

import statistics
import time

import dask
import dask.array as da
from dask.core import get_deps
from dask.utils import ensure_dict

try:
    from dask.base import collections_to_expr

    new_style = True
except ImportError:
    from dask.base import collections_to_dsk

    new_style = False


print(f"Dask version: {dask.__version__}")

niterations = 10


def benchmark(shape, chunks, indexer, identifier):
    try:
        durations = []
        for i in range(niterations):
            arr = da.ones(shape, chunks=chunks)
            start = time.perf_counter()
            new = arr[indexer]
            if new_style:
                dsk = collections_to_expr([new]).optimize().__dask_graph__()
            else:
                dsk = ensure_dict(collections_to_dsk([new]))
            get_deps(dsk)
            end = time.perf_counter()
            durations.append(end - start)
        print(
            f"{identifier:<40}: {sum(durations)/niterations * 1000:>10.2f} ms"
            # ± {statistics.stdev(durations):.2f} ms"
        )
    except Exception as e:
        print(f"Error with {identifier}: {e}")
        return


benchmark((100_000,), (1,), slice(None, None, None), "Empty indexer (all)")
benchmark((100_000,), (1,), slice(None, None, 2), "Every second element (::2)")
benchmark((100_000,), (1,), slice(None, -1), "All but last element (:-1)")


benchmark(
    (100, 1000, 100),
    (-1, 1, 1),
    (1, slice(None)),
    "Multi dim one slice (0, slice(None))",
)
benchmark(
    (100, 1000, 100),
    (-1, 1, 1),
    (1, slice(None, None, 2)),
    "Multi dim (0, slice(None, None, 2))",
)
benchmark(
    (100, 1000, 100), (-1, 1, 1), (1, slice(None, -1)), "Multi dim (0, slice(None, -1))"
)

benchmark(
    (1000, 35, 720, 1440, 31),
    (1, -1, 72, 72, -1),
    (slice(None, -1), slice(None, 10)),
    "Multi dim (slice(None, -1), slice(None, 10))",
)

```

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._

      9 files  ±0        9 suites  ±0   3h 17m 41s ⏱️ +19s
 18 032 tests ±0   16 816 ✅ ±0   1 216 💤 ±0  0 ❌ ±0 
161 362 runs  ±0  149 229 ✅ +3  12 133 💤  - 3  0 ❌ ±0 

Results for commit 0a0f974b. ± Comparison against base commit abdb435d.

[test-results]:data:application/gzip;base64,H4sIACmnJGgC/03MTQ6DIBCG4asY1l0wjA5ML9MglYTUnwZh1fTuRaPU3bzPJN9H+DAOq7g3fGvEmkOq8czRprDMJQEMQaHyTNsbjER19mPNzm1IBuiPr/DeUF3M2zAWkxWGGJd4SMzzPk2ApI6u0y0rxRXrNOJpl+m9r8tumaaQSghppWfd9kqzldQO2HdSA2omU050znqGTqH4/gCMDXd/GAEAAA==


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 9fd4f75a7..f751439d3 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -143,17 +143,13 @@ def parse_input(obj: Any) -> object:
         if any(isinstance(v, GraphNode) for v in parsed_dict.values()):
             return Dict(parsed_dict)
 
-    if isinstance(obj, (list, set, tuple)):
+    if isinstance(obj, (list, set)):
         parsed_collection = tuple(parse_input(o) for o in obj)
         if any(isinstance(o, GraphNode) for o in parsed_collection):
             if isinstance(obj, list):
                 return List(*parsed_collection)
             if isinstance(obj, set):
                 return Set(*parsed_collection)
-            if isinstance(obj, tuple):
-                if is_namedtuple_instance(obj):
-                    return _wrap_namedtuple_task(None, obj, parse_input)
-                return Tuple(*parsed_collection)
 
     return obj
 
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index 76419dc17..f456bacd3 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -325,31 +325,46 @@ def slice_slices_and_integers(
     sorted_block_slices = [sorted(i.items()) for i in block_slices]
 
     # (in_name, 1, 1, 2), (in_name, 1, 1, 4), (in_name, 2, 1, 2), ...
-    in_names = list(product([in_name], *[pluck(0, s) for s in sorted_block_slices]))
+    in_names = product([in_name], *[pluck(0, s) for s in sorted_block_slices])
 
     # (out_name, 0, 0, 0), (out_name, 0, 0, 1), (out_name, 0, 1, 0), ...
-    out_names = list(
-        product(
-            [out_name],
-            *[
-                range(len(d))[::-1] if i.step and i.step < 0 else range(len(d))
-                for d, i in zip(block_slices, index)
-                if not isinstance(i, Integral)
-            ],
-        )
+    out_names = product(
+        [out_name],
+        *[
+            range(len(d))[::-1] if i.step and i.step < 0 else range(len(d))
+            for d, i in zip(block_slices, index)
+            if not isinstance(i, Integral)
+        ],
     )
 
-    all_slices = list(product(*[pluck(1, s) for s in sorted_block_slices]))
-
-    dsk_out = {
-        out_name: (
-            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 Alias(out_name, in_name)
+    all_slices = product(*(pluck(1, s) for s in sorted_block_slices))
+    if not allow_getitem_optimization:
+        dsk_out = {
+            out_name: (Task(out_name, getitem, TaskRef(in_name), slices))
+            for out_name, in_name, slices in zip(out_names, in_names, all_slices)
+        }
+    else:
+        empty_slice = slice(None, None, None)
+
+        all_empty_slices = map(
+            all,
+            product(
+                *(
+                    (sp == empty_slice for sp in pluck(1, s))
+                    for s in sorted_block_slices
+                )
+            ),
         )
-        for out_name, in_name, slices in zip(out_names, in_names, all_slices)
-    }
+        dsk_out = {
+            out_name: (
+                Task(out_name, getitem, TaskRef(in_name), slices)
+                if not all_empty
+                else Alias(out_name, in_name)
+            )
+            for out_name, in_name, all_empty, slices in zip(
+                out_names, in_names, all_empty_slices, all_slices
+            )
+        }
 
     new_blockdims = [
         new_blockdim(d, db, i)
@@ -590,7 +605,14 @@ def take(outname, inname, chunks, index, axis=0):
 
         arange = arange_safe(np.sum(chunks[axis]), like=index)
         if len(index) == len(arange) and np.abs(index - arange).sum() == 0:
-            raise SlicingNoop()
+            # TODO: This should be a real no-op, but the call stack is
+            # too deep to do this efficiently for now
+            chunk_tuples = product(*(range(len(c)) for i, c in enumerate(chunks)))
+            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]))
 


Test patch:
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index daec4c1fd..b325830d7 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -259,6 +259,13 @@ def test_slice_array_2d():
     assert expected == result
 
 
+def test_mixed_index():
+    da_array = da.ones((1, 1, 31, 40))
+    new = da_array[(np.array([0]), 0, slice(None), slice(None))]
+    assert isinstance(new, da.Array)
+    assert_eq(new, da_array[0])
+
+
 def test_slice_optimizations():
     # bar[:]
     with pytest.raises(SlicingNoop):
@@ -327,8 +334,9 @@ def test_take_sorted():
     chunks, dsk = take("y-y", "x", [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
     assert len(dsk) == 6
     assert chunks == ((4,),)
-    with pytest.raises(SlicingNoop):
-        take("y", "x", [(20, 20, 20, 20)], np.arange(0, 80), axis=0)
+    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], np.arange(0, 80), axis=0)
+    assert len(dsk) == 4
+    assert chunks == ((20, 20, 20, 20),)
 
 
 def test_slicing_chunks():
diff --git a/dask/array/tests/test_xarray.py b/dask/array/tests/test_xarray.py
index f35e1bd42..e3c0d6f36 100644
--- a/dask/array/tests/test_xarray.py
+++ b/dask/array/tests/test_xarray.py
@@ -154,3 +154,13 @@ def test_shared_tasks(wrap_xarray):
     else:
         assert isinstance(comp_res[0], np.ndarray)
     assert total_calls == in1.blocks.size
+
+
+def test_slicing():
+    ds = xr.Dataset(
+        {"z": (("t", "p", "y", "x"), np.ones((1, 1, 31, 40)))},
+    )
+    ds = ds.chunk()
+    subset = ds.isel(t=[0], p=0).z[:, ::10, ::10][:, ::-1, :]
+    subset2 = ds.isel(t=[0]).isel(p=0).z[:, ::10, ::10][:, ::-1, :]
+    assert_eq(subset.data, subset2.data)
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 4986c3efb..36074e98e 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -499,32 +499,6 @@ def test_parse_graph_namedtuple_legacy(typ, args, kwargs):
     assert new_dsk["foo"]() == typ(*args, **kwargs)
 
 
-@pytest.mark.parametrize(
-    "typ, args, kwargs",
-    [
-        (PlainNamedTuple, ["some-data"], {}),
-        (NewArgsNamedTuple, ["some", "data", "more"], {}),
-        (NewArgsExNamedTuple, ["some", "data", "more"], {"another": "data"}),
-    ],
-)
-def test_parse_namedtuple(typ, args, kwargs):
-    def func(x):
-        return x
-
-    obj = typ(*args, **kwargs)
-    t = Task("foo", func, parse_input(obj))
-
-    assert t() == obj
-
-    # The other test tuple do weird things to their input
-    if typ is PlainNamedTuple:
-        args = tuple([TaskRef("b")] + list(args)[1:])
-        obj = typ(*args, **kwargs)
-        t = Task("foo", func, parse_input(obj))
-        assert t.dependencies == {"b"}
-        assert t({"b": "foo"}) == typ(*tuple(["foo"] + list(args)[1:]), **kwargs)
-
-
 def test_pickle_literals():
     np = pytest.importorskip("numpy")
     obj = DataNode("foo", np.transpose)