Repository: dask/dask

Issue:
Ensure that array-shuffle with range-like full indexer is a no-op
Assume

```
arr = da.from_array(...)

arr.shuffle(indexer)
```

where the flattened indexer is basically

```
list(range(0, len(are))
```

where the sub-lists also match the chunk structure
This should be a no-op. This is unlikely to be triggered in Dask, but xarray adds a shuffle as well and ``arr.shuffle().shuffle()`` should only trigger one shuffle there.

See https://github.com/pydata/xarray/pull/9320#discussion_r1707501882 for context

Discussion:
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   2h 56m 35s :stopwatch: + 2m 59s
 13 133 tests + 2   12 065 :white_check_mark: + 2   1 068 :zzz: ±0  0 :x: ±0 
136 849 runs  +22  117 784 :white_check_mark: +21  19 065 :zzz: +1  0 :x: ±0 

Results for commit 353873ea. ± Comparison against base commit ed5f6889.

[test-results]:data:application/gzip;base64,H4sIADn3tWYC/02NSQ7DIBAEv2JxzgE8DEs+E5nFEoqXCMPJyt+Dl2Afq1pTs5I+DH4hz4bBoyFLDukil2OXwjxtTFFjUWVNx84A/vxasrWbbKnAS77DZ78Uqrq+C0NxtAof4xxPE/N0pIXi+uSaZlIqXuWZ1se/3d3SO9/Ldh7HkAoQQFASfGeM544rJrSkRmvRMjQOmTNgqOUWyfcHt56uZRoBAAA=


Patch:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 89f435d7b..6e3dd2e4a 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -71,6 +71,9 @@ def shuffle(x, indexer: list[list[int]], axis):
     out_name = f"shuffle-{token}"
 
     chunks, layer = _shuffle(x.chunks, indexer, axis, x.name, out_name, token)
+    if len(layer) == 0:
+        return Array(x.dask, x.name, x.chunks, meta=x)
+
     graph = HighLevelGraph.from_collections(out_name, layer, dependencies=[x])
 
     return Array(graph, out_name, chunks, meta=x)
@@ -89,6 +92,17 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
         raise IndexError(
             f"Indexer contains out of bounds index. Dimension only has {sum(chunks[axis])} elements."
         )
+
+    if len(indexer) == len(chunks[axis]):
+        # check if the array is already shuffled the way we want
+        ctr = 0
+        for idx, c in zip(indexer, chunks[axis]):
+            if idx != list(range(ctr, ctr + c)):
+                break
+            ctr += c
+        else:
+            return chunks, {}
+
     indexer = copy.deepcopy(indexer)
 
     chunksize_tolerance = config.get("array.shuffle.chunksize-tolerance")
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index 3f345cfb8..0a0421fd7 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -566,6 +566,14 @@ def take(outname, inname, chunks, index, axis=0):
     if not np.isnan(chunks[axis]).any():
         from dask.array._shuffle import _shuffle
 
+        arange = np.arange(np.sum(chunks[axis]))
+        if len(index) == len(arange) and np.abs(index - arange).sum() == 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}
+            return tuple(chunks), graph
+
         average_chunk_size = int(sum(chunks[axis]) / len(chunks[axis]))
 
         indexer = []


Test patch:
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 469f22be0..2882cbc87 100644
--- a/dask/array/tests/test_shuffle.py
+++ b/dask/array/tests/test_shuffle.py
@@ -81,3 +81,16 @@ def test_oob_axis(darr):
 def test_oob_indexer(darr):
     with pytest.raises(IndexError, match="Indexer contains out of bounds index"):
         darr.shuffle([[16]], axis=1)
+
+
+def test_shuffle_no_op_with_correct_indexer():
+    arr = da.ones((250, 100), chunks=((50, 100, 33, 67), 100))
+    indexer = [
+        list(range(0, 50)),
+        list(range(50, 150)),
+        list(range(150, 183)),
+        list(range(183, 250)),
+    ]
+    result = arr.shuffle(indexer, axis=0)
+    assert result.dask == arr.dask
+    assert_eq(arr, result)
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index 0e9cfce28..f503c5fee 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -413,10 +413,10 @@ def test_take_sorted():
 
     chunks, dsk = take("y", "x", [(20, 20, 20, 20)], np.arange(0, 80), axis=0)
     expected = {
-        ("y", 0): (getitem, ("x", 0), (np.arange(0, 20),)),
-        ("y", 1): (getitem, ("x", 1), (np.arange(0, 20),)),
-        ("y", 2): (getitem, ("x", 2), (np.arange(0, 20),)),
-        ("y", 3): (getitem, ("x", 3), (np.arange(0, 20),)),
+        ("y", 0): ("x", 0),
+        ("y", 1): ("x", 1),
+        ("y", 2): ("x", 2),
+        ("y", 3): ("x", 3),
     }
     np.testing.assert_equal(dsk, expected)
     assert chunks == ((20, 20, 20, 20),)
@@ -1102,3 +1102,20 @@ def test_slice_array_null_dimension():
     array = da.from_array(np.zeros((3, 0)))
     expected = np.zeros((3, 0))[[0]]
     assert_eq(array[[0]], expected)
+
+
+def test_take_sorted_indexer():
+    arr = da.ones((250, 100), chunks=((50, 100, 33, 67), 100))
+    indexer = list(range(0, 250))
+    result = arr[indexer, :]
+    assert_eq(arr, result)
+    assert {
+        **dict(arr.dask),
+        **{
+            k: k2
+            for k, k2 in zip(
+                [k for k in dict(result.dask) if "getitem" in k[0]],
+                dict(arr.dask).keys(),
+            )
+        },
+    } == dict(result.dask)