Repository: dask/dask
Pull request #11538 (2024-11-19 17:15:29)

Remove optimize_slices
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

This can all go away now

``fuse_slices`` should probably go away to or get a lot easier, but this is still used in store

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 17m 29s ⏱️ - 6m 42s
 13 255 tests  -  42   12 193 ✅  -  42   1 062 💤 ±0  0 ❌ ±0 
138 316 runs   - 462  119 321 ✅  - 466  18 995 💤 +4  0 ❌ ±0 

Results for commit 596f8af8. ± Comparison against base commit 5b115c43.

<details>
  <summary>This pull request <b>removes</b> 42 tests.</summary>

```
dask.array.tests.test_optimization ‑ test_dont_fuse_fancy_indexing_in_getter_nofancy[False-False]
dask.array.tests.test_optimization ‑ test_dont_fuse_fancy_indexing_in_getter_nofancy[False-True]
dask.array.tests.test_optimization ‑ test_dont_fuse_fancy_indexing_in_getter_nofancy[True-False]
dask.array.tests.test_optimization ‑ test_dont_fuse_fancy_indexing_in_getter_nofancy[True-True]
dask.array.tests.test_optimization ‑ test_dont_remove_no_op_slices_for_getter_or_getter_nofancy[False-False-False-getter]
dask.array.tests.test_optimization ‑ test_dont_remove_no_op_slices_for_getter_or_getter_nofancy[False-False-False-getter_nofancy]
dask.array.tests.test_optimization ‑ test_dont_remove_no_op_slices_for_getter_or_getter_nofancy[False-False-True-getter]
dask.array.tests.test_optimization ‑ test_dont_remove_no_op_slices_for_getter_or_getter_nofancy[False-False-True-getter_nofancy]
dask.array.tests.test_optimization ‑ test_dont_remove_no_op_slices_for_getter_or_getter_nofancy[False-True-False-getter]
dask.array.tests.test_optimization ‑ test_dont_remove_no_op_slices_for_getter_or_getter_nofancy[False-True-False-getter_nofancy]
…
```
</details>

[test-results]:data:application/gzip;base64,H4sIAKrPPGcC/03MSw6DIBSF4a0Yxh1wQW6hm2koj4TUR4MwMt17oSo6PN9J/pX4MLiFPDrgt44sOaRz2Rx1CvNUN8heFSpv2n4mxLGfSzamIgPFT3yHT0WKrJnXYShGG7gY57hLzNOWlhxw3y1dygwa7mmplDjskv7va9nM4xhSGUQo9FJ7aZBrROuNRY3UuztI/qLCIApne3Dk+wMAsUuaGgEAAA==


Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 10736af6f..0a7ebeb1b 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -3754,9 +3754,9 @@ def from_zarr(
                 store = zarr.storage.FSStore(url, **storage_options)
         else:
             store = url
-        z = zarr.open_array(store=store, read_only=True, path=component, **kwargs)
+        z = zarr.open_array(store=store, path=component, **kwargs)
     else:
-        z = zarr.open_array(store=url, read_only=True, path=component, **kwargs)
+        z = zarr.open_array(store=url, path=component, **kwargs)
     chunks = chunks if chunks is not None else z.chunks
     if name is None:
         name = "from-zarr-" + tokenize(z, component, storage_options, chunks, **kwargs)
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
index 3ec56353a..20893e30a 100644
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -1,36 +1,21 @@
 from __future__ import annotations
 
-from collections.abc import Callable
 from itertools import zip_longest
 from numbers import Integral
-from typing import Any
-
-import numpy as np
 
 from dask import config
 from dask._task_spec import convert_legacy_graph, fuse_linear_task_spec
-from dask.array.chunk import getitem
-from dask.array.core import getter, getter_inline, getter_nofancy
 from dask.blockwise import fuse_roots, optimize_blockwise
 from dask.core import flatten
 from dask.highlevelgraph import HighLevelGraph
-from dask.optimization import SubgraphCallable
 from dask.utils import ensure_dict
 
-# All get* functions the optimizations know about
-GETTERS = (getter, getter_nofancy, getter_inline, getitem)
-# These get* functions aren't ever completely removed from the graph,
-# even if the index should be a no-op by numpy semantics. Some array-like's
-# don't completely follow semantics, making indexing always necessary.
-GETNOREMOVE = (getter, getter_nofancy)
-
 
 def optimize(dsk, keys, **kwargs):
     """Optimize dask for array computation
 
     1.  Cull tasks not necessary to evaluate keys
-    2.  Remove full slicing, e.g. x[:]
-    3.  Inline fast functions like getitem and np.transpose
+    2.  Perform linear fusion
     """
     if not isinstance(keys, (list, set)):
         keys = [keys]
@@ -56,126 +41,6 @@ def optimize(dsk, keys, **kwargs):
     return dsk
 
 
-def _is_getter_task(
-    value,
-) -> tuple[Callable, Any, Any, bool, bool | None] | None:
-    """Check if a value in a Dask graph looks like a getter.
-
-    1. Is it a tuple with the first element a known getter.
-    2. Is it a SubgraphCallable with a single element in its
-       dsk which is a known getter.
-
-    If a getter is found, it returns a tuple with (getter, array, index, asarray, lock).
-    Otherwise it returns ``None``.
-
-    TODO: the second check is a hack to allow for slice fusion between tasks produced
-    from blockwise layers and slicing operations. Once slicing operations have
-    HighLevelGraph layers which can talk to Blockwise layers this check *should* be
-    removed, and we should not have to introspect SubgraphCallables.
-    """
-    if type(value) is not tuple:
-        return None
-    first = value[0]
-    get: Callable | None = None
-    if first in GETTERS:
-        get = first
-    # We only accept SubgraphCallables with a single sub-task right now as it's
-    # not clear which task to inspect if there is more than one, or how to resolve
-    # conflicts if they occur.
-    elif isinstance(first, SubgraphCallable) and len(first.dsk) == 1:
-        v = next(iter(first.dsk.values()))
-        if type(v) is tuple and len(v) > 1 and v[0] in GETTERS:
-            get = v[0]
-    if get is None:  # Didn't find a getter
-        return None
-
-    length = len(value)
-    if length == 3:
-        # getter defaults to asarray=True, getitem is semantically False
-        return get, value[1], value[2], get is not getitem, None
-    elif length == 5:
-        return get, *value[1:]
-
-    return None
-
-
-def optimize_slices(dsk):
-    """Optimize slices
-
-    1.  Fuse repeated slices, like x[5:][2:6] -> x[7:11]
-    2.  Remove full slices, like         x[:] -> x
-
-    See also:
-        fuse_slice_dict
-    """
-    fancy_ind_types = (list, np.ndarray)
-    dsk = dsk.copy()
-    for k, v in dsk.items():
-        if a_task := _is_getter_task(v):
-            get, a, a_index, a_asarray, a_lock = a_task
-
-            while b_task := _is_getter_task(a):
-                f2, b, b_index, b_asarray, b_lock = b_task
-
-                if a_lock and a_lock is not b_lock:
-                    break
-                if (type(a_index) is tuple) != (type(b_index) is tuple):
-                    break
-                if type(a_index) is tuple:
-                    indices = b_index + a_index
-                    if len(a_index) != len(b_index) and any(i is None for i in indices):
-                        break
-                    if f2 is getter_nofancy and any(
-                        isinstance(i, fancy_ind_types) for i in indices
-                    ):
-                        break
-                elif f2 is getter_nofancy and (
-                    type(a_index) in fancy_ind_types or type(b_index) in fancy_ind_types
-                ):
-                    break
-                try:
-                    c_index = fuse_slice(b_index, a_index)
-                    # rely on fact that nested gets never decrease in
-                    # strictness e.g. `(getter_nofancy, (getter, ...))` never
-                    # happens
-                    get = getter if f2 is getter_inline else f2
-                except NotImplementedError:
-                    break
-                a, a_index, a_lock = b, c_index, b_lock
-                a_asarray |= b_asarray
-
-            # Skip the get call if not from from_array and nothing to do
-            if get not in GETNOREMOVE and (
-                (
-                    type(a_index) is slice
-                    and not a_index.start
-                    and a_index.stop is None
-                    and a_index.step is None
-                )
-                or (
-                    type(a_index) is tuple
-                    and all(
-                        type(s) is slice
-                        and not s.start
-                        and s.stop is None
-                        and s.step is None
-                        for s in a_index
-                    )
-                )
-            ):
-                dsk[k] = a
-            elif get is getitem or (a_asarray and not a_lock):
-                # default settings are fine, drop the extra parameters Since we
-                # always fallback to inner `get` functions, `get is getitem`
-                # can only occur if all gets are getitem, meaning all
-                # parameters must be getitem defaults.
-                dsk[k] = (get, a, a_index)
-            else:
-                dsk[k] = (get, a, a_index, a_asarray, a_lock)
-
-    return dsk
-
-
 def normalize_slice(s):
     """Replace Nones in slices with integers
 
diff --git a/dask/array/reshape.py b/dask/array/reshape.py
index 1cc79090a..b824110c7 100644
--- a/dask/array/reshape.py
+++ b/dask/array/reshape.py
@@ -473,6 +473,9 @@ def reshape_blockwise(
 
     _sanity_checks(x, shape)
 
+    if len(shape) == x.ndim and shape == x.shape:
+        return Array(x.dask, x.name, x.chunks, meta=x)
+
     outname = "reshape-blockwise-" + tokenize(x, shape)
     chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(x.chunks))))
 
@@ -515,9 +518,6 @@ def reshape_blockwise(
         x.shape, shape, x.chunks, disallow_dimension_expansion=True
     )
 
-    if len(shape) == x.ndim:
-        return Array(x.dask, x.name, x.chunks, meta=x)
-
     # Convert input chunks to output chunks
     out_shapes = [
         _convert_to_shape(c, mapper_in, one_dimensions)


Test patch:
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
index cf8583bb5..66d25551d 100644
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -13,16 +13,10 @@ import dask.array as da
 from dask.array.chunk import getitem as da_getitem
 from dask.array.core import getter as da_getter
 from dask.array.core import getter_nofancy as da_getter_nofancy
-from dask.array.optimization import (
-    fuse_slice,
-    optimize,
-    optimize_blockwise,
-    optimize_slices,
-)
+from dask.array.optimization import fuse_slice, optimize, optimize_blockwise
 from dask.array.utils import assert_eq
 from dask.highlevelgraph import HighLevelGraph
-from dask.optimization import SubgraphCallable, fuse
-from dask.utils import SerializableLock
+from dask.optimization import SubgraphCallable
 
 
 def _wrap_getter(func, wrap):
@@ -108,148 +102,6 @@ def _assert_getter_dsk_eq(a, b):
             assert av == bv
 
 
-def test_fuse_getitem(getter, getter_nofancy, getitem):
-    pairs = [
-        (
-            (getter, (getter, "x", slice(1000, 2000)), slice(15, 20)),
-            (getter, "x", slice(1015, 1020)),
-        ),
-        (
-            (
-                getitem,
-                (getter, "x", (slice(1000, 2000), slice(100, 200))),
-                (slice(15, 20), slice(50, 60)),
-            ),
-            (getter, "x", (slice(1015, 1020), slice(150, 160))),
-        ),
-        (
-            (
-                getitem,
-                (getter_nofancy, "x", (slice(1000, 2000), slice(100, 200))),
-                (slice(15, 20), slice(50, 60)),
-            ),
-            (getter_nofancy, "x", (slice(1015, 1020), slice(150, 160))),
-        ),
-        ((getter, (getter, "x", slice(1000, 2000)), 10), (getter, "x", 1010)),
-        (
-            (getitem, (getter, "x", (slice(1000, 2000), 10)), (slice(15, 20),)),
-            (getter, "x", (slice(1015, 1020), 10)),
-        ),
-        (
-            (getitem, (getter_nofancy, "x", (slice(1000, 2000), 10)), (slice(15, 20),)),
-            (getter_nofancy, "x", (slice(1015, 1020), 10)),
-        ),
-        (
-            (getter, (getter, "x", (10, slice(1000, 2000))), (slice(15, 20),)),
-            (getter, "x", (10, slice(1015, 1020))),
-        ),
-        (
-            (
-                getter,
-                (getter, "x", (slice(1000, 2000), slice(100, 200))),
-                (slice(None, None), slice(50, 60)),
-            ),
-            (getter, "x", (slice(1000, 2000), slice(150, 160))),
-        ),
-        (
-            (getter, (getter, "x", (None, slice(None, None))), (slice(None, None), 5)),
-            (getter, "x", (None, 5)),
-        ),
-        (
-            (
-                getter,
-                (getter, "x", (slice(1000, 2000), slice(10, 20))),
-                (slice(5, 10),),
-            ),
-            (getter, "x", (slice(1005, 1010), slice(10, 20))),
-        ),
-        (
-            (
-                getitem,
-                (getitem, "x", (slice(1000, 2000),)),
-                (slice(5, 10), slice(10, 20)),
-            ),
-            (getitem, "x", (slice(1005, 1010), slice(10, 20))),
-        ),
-        (
-            (getter, (getter, "x", slice(1000, 2000), False, False), slice(15, 20)),
-            (getter, "x", slice(1015, 1020)),
-        ),
-        (
-            (getter, (getter, "x", slice(1000, 2000)), slice(15, 20), False, False),
-            (getter, "x", slice(1015, 1020)),
-        ),
-        (
-            (
-                getter,
-                (getter_nofancy, "x", slice(1000, 2000), False, False),
-                slice(15, 20),
-                False,
-                False,
-            ),
-            (getter_nofancy, "x", slice(1015, 1020), False, False),
-        ),
-    ]
-
-    for inp, expected in pairs:
-        result = optimize_slices({"y": inp})
-        _assert_getter_dsk_eq(result, {"y": expected})
-
-
-def test_fuse_getitem_lock(getter, getter_nofancy, getitem):
-    lock1 = SerializableLock()
-    lock2 = SerializableLock()
-
-    pairs = [
-        (
-            (getter, (getter, "x", slice(1000, 2000), True, lock1), slice(15, 20)),
-            (getter, "x", slice(1015, 1020), True, lock1),
-        ),
-        (
-            (
-                getitem,
-                (getter, "x", (slice(1000, 2000), slice(100, 200)), True, lock1),
-                (slice(15, 20), slice(50, 60)),
-            ),
-            (getter, "x", (slice(1015, 1020), slice(150, 160)), True, lock1),
-        ),
-        (
-            (
-                getitem,
-                (
-                    getter_nofancy,
-                    "x",
-                    (slice(1000, 2000), slice(100, 200)),
-                    True,
-                    lock1,
-                ),
-                (slice(15, 20), slice(50, 60)),
-            ),
-            (getter_nofancy, "x", (slice(1015, 1020), slice(150, 160)), True, lock1),
-        ),
-        (
-            (
-                getter,
-                (getter, "x", slice(1000, 2000), True, lock1),
-                slice(15, 20),
-                True,
-                lock2,
-            ),
-            (
-                getter,
-                (getter, "x", slice(1000, 2000), True, lock1),
-                slice(15, 20),
-                True,
-                lock2,
-            ),
-        ),
-    ]
-
-    for inp, expected in pairs:
-        result = optimize_slices({"y": inp})
-        _assert_getter_dsk_eq(result, {"y": expected})
-
-
 def test_optimize_with_getitem_fusion(getter):
     dsk = {
         "a": "some-array",
@@ -262,30 +114,6 @@ def test_optimize_with_getitem_fusion(getter):
     assert len(result) < len(dsk)
 
 
-def test_optimize_slicing(getter):
-    dsk = {
-        "a": (range, 10),
-        "b": (getter, "a", (slice(None, None, None),)),
-        "c": (getter, "b", (slice(None, None, None),)),
-        "d": (getter, "c", (slice(0, 5, None),)),
-        "e": (getter, "d", (slice(None, None, None),)),
-    }
-
-    expected = {"e": (getter, (range, 10), (slice(0, 5, None),))}
-    result = optimize_slices(fuse(dsk, [], rename_keys=False)[0])
-    _assert_getter_dsk_eq(result, expected)
-
-    # protect output keys
-    expected = {
-        "c": (getter, (range, 10), (slice(0, None, None),)),
-        "d": (getter, "c", (slice(0, 5, None),)),
-        "e": (getter, "d", (slice(None, None, None),)),
-    }
-    result = optimize_slices(fuse(dsk, ["c", "d", "e"], rename_keys=False)[0])
-
-    _assert_getter_dsk_eq(result, expected)
-
-
 def test_fuse_slice():
     assert fuse_slice(slice(10, 15), slice(0, 5, 2)) == slice(10, 15, 2)
 
@@ -342,13 +170,6 @@ def test_nonfusible_fancy_indexing():
             fuse_slice(a, b)
 
 
-def test_hard_fuse_slice_cases(getter):
-    dsk = {
-        "x": (getter, (getter, "x", (None, slice(None, None))), (slice(None, None), 5))
-    }
-    _assert_getter_dsk_eq(optimize_slices(dsk), {"x": (getter, "x", (None, 5))})
-
-
 def test_dont_fuse_numpy_arrays():
     x = np.ones(10)
     for _ in [(5,), (10,)]:
@@ -376,20 +197,6 @@ def test_fuse_slices_with_alias(getter, getitem):
     assert len(dsk2) == 3
 
 
-def test_dont_fuse_fancy_indexing_in_getter_nofancy(getitem, getter_nofancy):
-    dsk = {
-        "a": (
-            getitem,
-            (getter_nofancy, "x", (slice(10, 20, None), slice(100, 200, None))),
-            ([1, 3], slice(50, 60, None)),
-        )
-    }
-    _assert_getter_dsk_eq(optimize_slices(dsk), dsk)
-
-    dsk = {"a": (getitem, (getter_nofancy, "x", [1, 2, 3]), 0)}
-    _assert_getter_dsk_eq(optimize_slices(dsk), dsk)
-
-
 @pytest.mark.parametrize("chunks", [10, 5, 3])
 def test_fuse_getter_with_asarray(chunks):
     x = np.ones(10) * 1234567890
@@ -403,50 +210,6 @@ def test_fuse_getter_with_asarray(chunks):
     assert_eq(z, x + 1)
 
 
-def test_remove_no_op_slices_for_getitem(getitem):
-    null = slice(0, None)
-    opts = [
-        ((getitem, "x", null, False, False), "x"),
-        ((getitem, (getitem, "x", null, False, False), null), "x"),
-        ((getitem, (getitem, "x", (null, null), False, False), ()), "x"),
-    ]
-    for orig, final in opts:
-        _assert_getter_dsk_eq(optimize_slices({"a": orig}), {"a": final})
-
-
-@pytest.mark.parametrize("which", ["getter", "getter_nofancy"])
-def test_dont_remove_no_op_slices_for_getter_or_getter_nofancy(
-    which, getitem, getter, getter_nofancy
-):
-    # Test that no-op slices are *not* removed when using getter or
-    # getter_nofancy. This ensures that `get` calls are always made in all
-    # tasks created by `from_array`, even after optimization
-
-    # Pytest doesn't make it easy to parameterize over parameterized fixtures
-    if which == "getter":
-        get = getter
-    else:
-        get = getter_nofancy
-
-    null = slice(0, None)
-    opts = [
-        (
-            (get, "x", null, False, False),
-            (get, "x", null, False, False),
-        ),
-        (
-            (getitem, (get, "x", null, False, False), null),
-            (get, "x", null, False, False),
-        ),
-        (
-            (getitem, (get, "x", (null, null), False, False), ()),
-            (get, "x", (null, null), False, False),
-        ),
-    ]
-    for orig, final in opts:
-        _assert_getter_dsk_eq(optimize_slices({"a": orig}), {"a": final})
-
-
 @pytest.mark.xfail(reason="blockwise fusion does not respect this, which is ok")
 def test_turn_off_fusion():
     x = da.ones(10, chunks=(5,))
diff --git a/dask/array/tests/test_reshape.py b/dask/array/tests/test_reshape.py
index 14fb84ba9..1473d97d1 100644
--- a/dask/array/tests/test_reshape.py
+++ b/dask/array/tests/test_reshape.py
@@ -411,12 +411,11 @@ def test_reshape_blockwise_raises_for_expansion():
         reshape_blockwise(arr, (3, 2, 9))
 
 
-@pytest.mark.parametrize("shape", [(18, 3), (6, 3, 3)])
-def test_reshape_blockwise_dimension_reduction_and_chunks(shape):
+def test_reshape_blockwise_dimension_reduction_and_chunks():
     x = np.arange(0, 54).reshape(6, 3, 3)
     arr = from_array(x, chunks=(3, 2, (2, 1)))
     with pytest.raises(ValueError, match="Setting chunks is not allowed when"):
-        reshape_blockwise(arr, shape, arr.chunks)
+        reshape_blockwise(arr, (18, 3), arr.chunks)
 
 
 def test_reshape_blockwise_identity():
@@ -426,6 +425,10 @@ def test_reshape_blockwise_identity():
     assert len(arr.dask) == len(result.dask)
     assert_eq(result, x)
 
+    result = reshape_blockwise(arr, (6, 3, 3), chunks=arr.chunks)
+    assert len(arr.dask) == len(result.dask)
+    assert_eq(result, x)
+
 
 def test_argwhere_reshaping_not_concating_chunks():
     # GH#10080