Repository: dask/dask
Pull request #11364 (2024-09-02 20:07:23)

Cast indexer to minimal dtype in shuffle
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

from 55MiB to 21MiB:

```

arr = da.ones((3_000_000, 16, 16), chunks=(10_000, 16, 16))
arr = arr.persist()

indexer = np.arange(3_000_000)
np.random.shuffle(indexer)

result = arr[indexer, :, :]
result.persist()
```

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 2m 28s :stopwatch: + 2m 2s
 13 170 tests + 1   12 114 :white_check_mark: + 1   1 056 :zzz: ±0  0 :x: ±0 
137 257 runs  +11  118 230 :white_check_mark: +10  19 027 :zzz: +1  0 :x: ±0 

Results for commit e33132aa. ± Comparison against base commit 2cfb3ef0.

[test-results]:data:application/gzip;base64,H4sIAKUi1mYC/02MSw7DIAwFrxKx7gJDSEgvUzmukVDzqQisqt695FOS5Yz95iOcH3gR9wr0rRJL8vGkZwoY/TytLLvaZpWvcb9DK//8WBLRKhVAfcqXf29L0xTn0A/ZnUsOYQ6HCWna060y7cElDVZpWeSR7qQqj5f0xtcyzePoYwbBWoNWiBaJEB1rapzpjea6I2ZG01s0DFJ8fxVz3DUaAQAA


Patch:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index fd0a51616..25320100d 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -207,6 +207,7 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
 
     intermediates = dict()
     merges = dict()
+    dtype = _minimal_dtype(max(chunks[axis]))
     split_name = f"shuffle-split-{token}"
     slices = [slice(None)] * len(chunks)
     split_name_suffixes = count()
@@ -217,7 +218,7 @@ 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)
+        sorter = np.argsort(new_chunk_taker).astype(dtype)
         sorted_array = new_chunk_taker[sorter]
         source_chunk_nr, taker_boundary = np.unique(
             np.searchsorted(chunk_boundaries, sorted_array, side="right"),
@@ -243,9 +244,10 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
                 if c in taker_cache:
                     this_slice[axis] = taker_cache[c]
                 else:
-                    this_slice[axis] = sorted_array[b_start:b_end] - (
-                        chunk_boundaries[c - 1] if c > 0 else 0
-                    )
+                    this_slice[axis] = (
+                        sorted_array[b_start:b_end]
+                        - (chunk_boundaries[c - 1] if c > 0 else 0)
+                    ).astype(dtype)
                     if len(source_chunk_nr) == 1:
                         this_slice[axis] = this_slice[axis][np.argsort(sorter)]
                     taker_cache[c] = this_slice[axis]
@@ -286,3 +288,12 @@ def convert_key(key, chunk, axis):
     key = list(key)
     key.insert(axis, chunk)
     return tuple(key)
+
+
+def _minimal_dtype(max_value):
+    for dtype in [np.uint8, np.uint16, np.uint32, np.uint64]:
+        if max_value <= np.iinfo(dtype).max:
+            return dtype
+    else:
+        # shouldn't happen
+        raise OverflowError(f"Can't find a dtype for {max_value}")
diff --git a/dask/order.py b/dask/order.py
index 9ee946805..19e85a056 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -600,7 +600,7 @@ def order(
 
 def _connecting_to_roots(
     dependencies: Mapping[Key, set[Key]], dependents: Mapping[Key, set[Key]]
-) -> tuple[dict[Key, set[Key]], dict[Key, int]]:
+) -> tuple[dict[Key, frozenset[Key]], dict[Key, int]]:
     """Determine for every node which root nodes are connected to it (i.e.
     ancestors). If arguments of dependencies and dependents are switched, this
     can also be used to determine which leaf nodes are connected to which node
@@ -625,47 +625,59 @@ def _connecting_to_roots(
             # benefit from the speedup of using integers would be to convert
             # this back on demand which makes the code very hard to read.
             roots.add(k)
-            result[k] = {k}
+            result[k] = frozenset({k})
             deps = dependents[k]
             max_dependents[k] = len(deps)
             for child in deps:
                 num_needed[child] -= 1
                 if not num_needed[child]:
                     current.append(child)
+
+    dedup_mapping: dict[frozenset[Key], frozenset[Key]] = {}
     while current:
         key = current.pop()
+        if key in result:
+            continue
         for parent in dependents[key]:
             num_needed[parent] -= 1
             if not num_needed[parent]:
                 current.append(parent)
         # At some point, all the roots are the same, particularly for dense
         # graphs. We don't want to create new sets over and over again
-        new_set = None
+        new_set: set | None = None
         identical_sets = True
-        result_first = None
+        result_first: frozenset | None = None
 
-        for child in dependencies[key]:
+        for child in sorted(
+            dependencies[key], key=lambda k: len(dependents[k]), reverse=True
+        ):
             r_child = result[child]
-            if not result_first:
+            if result_first is None:
                 result_first = r_child
                 max_dependents[key] = max_dependents[child]
             # This clause is written such that it can circuit break early
-            elif not (  # type: ignore[unreachable]
+            elif not (
                 identical_sets
                 and (result_first is r_child or r_child.issubset(result_first))
             ):
                 identical_sets = False
                 if not new_set:
-                    new_set = result_first.copy()
+                    new_set = set(result_first)
                 max_dependents[key] = max(max_dependents[child], max_dependents[key])
                 new_set.update(r_child)
 
-        assert new_set is not None or result_first is not None
-        result[key] = new_set or result_first
+        if new_set:
+            new_set_frozen = frozenset(new_set)
+            deduped = dedup_mapping.get(new_set_frozen, None)
+            if deduped is None:
+                dedup_mapping[new_set_frozen] = deduped = new_set_frozen
+            result[key] = deduped
+        else:
+            assert result_first is not None
+            result[key] = result_first
+    del dedup_mapping
 
-    # The order algo doesn't care about this but this makes it easier to
-    # understand and shouldn't take that much time
-    empty_set: set[Key] = set()
+    empty_set: frozenset[Key] = frozenset()
     for r in roots:
         result[r] = empty_set
     return result, max_dependents


Test patch:
diff --git a/.github/workflows/test-report.yaml b/.github/workflows/test-report.yaml
index 1fb49605b..01d09e292 100644
--- a/.github/workflows/test-report.yaml
+++ b/.github/workflows/test-report.yaml
@@ -61,7 +61,7 @@ jobs:
           mv test_report.html test_short_report.html deploy/
 
       - name: Deploy 🚀
-        uses: JamesIves/github-pages-deploy-action@v4.6.3
+        uses: JamesIves/github-pages-deploy-action@v4.6.4
         with:
           branch: gh-pages
           folder: deploy
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 3458a443d..70339768b 100644
--- a/dask/array/tests/test_shuffle.py
+++ b/dask/array/tests/test_shuffle.py
@@ -140,3 +140,15 @@ def test_resize_other_dimensions():
     result = _rechunk_other_dimensions(arr, 4, 1, "auto")
     assert result.chunks == ((1, 1), (1,), (1, 1))
     assert_eq(arr, result)
+
+
+def test_dtype_taker(arr, darr):
+    indexer = [[1, 5, 6], [0, 3], [4, 2, 7]]
+    result = darr.shuffle(indexer, axis=1)
+    expected = arr[:, [1, 5, 6, 0, 3, 4, 2, 7]]
+    assert_eq(result, expected)
+    assert all(
+        v[1].dtype == np.uint8
+        for k, v in dict(result.dask).items()
+        if "array-taker" in k
+    )
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index bea8caf73..008795b2f 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+from collections import defaultdict
 
 import pytest
 
@@ -2401,6 +2402,7 @@ def test_connecting_to_roots_tree_reduction():
         "c2": {"a3", "a2"},
         "d": {"a1", "a2", "a3", "a0"},
     }
+    assert len({id(v) for v in connected_roots.values() if len(v) == 0}) == 1
     assert all(v == 1 for v in max_dependents.values())
 
     connected_roots, max_dependents = _connecting_to_roots(dependents, dependencies)
@@ -2411,6 +2413,78 @@ def test_connecting_to_roots_tree_reduction():
     assert all(v == 2 for v in max_dependents.values())
 
 
+def test_connecting_to_roots_long_linear_chains():
+    dsk = {
+        "a0": (f, 1),
+        "a1": (f, 1),
+        "b0": (f, "a0"),
+        "c0": (f, "b0"),
+        "d0": (f, "c0"),
+        "d1": (f, "c0"),
+        "e1": (f, "d0", "d1"),
+        "b1": (f, "a1"),
+        "c1": (f, "b1"),
+        "d2": (f, "c1"),
+        "d3": (f, "c1"),
+        "f": (f, "e1", "d2", "d3", "d1"),
+    }
+    dependencies, dependents = get_deps(dsk)
+    connected_roots, _ = _connecting_to_roots(dependencies, dependents)
+
+    id_a0 = id(connected_roots["b0"])
+    for k in ["c0", "d0", "d1", "e1"]:
+        assert id(connected_roots[k]) == id_a0
+
+    id_a0 = id(connected_roots["b1"])
+    for k in ["c1", "d2", "d3"]:
+        assert id(connected_roots[k]) == id_a0
+
+
+def test_connected_roots_deduplication_mem_usage():
+    # see https://github.com/dask/dask/issues/11055
+    np = pytest.importorskip("numpy")
+    da = pytest.importorskip("dask.array")
+
+    def rotate(X, n_iter=2):
+        gamma = 1
+        n_samples, n_modes = X.shape
+        R = np.eye(n_modes)
+
+        for _ in range(n_iter):
+            basis = X @ R
+
+            basis2 = basis * basis.conj()
+            basis3 = basis2 * basis
+            W = np.diag(np.sum(basis2, axis=0))
+            alpha = gamma / n_samples
+
+            transformed = X.conj().T @ (basis3 - (alpha * basis @ W))
+            U, svals, VT = da.linalg.svd_compressed(
+                transformed, k=n_modes, compute=False
+            )
+            R = U @ VT
+
+        Xrot = X @ R
+        return Xrot
+
+    n_modes = 100
+
+    x = da.random.random((10_000, 100), chunks=(5_000, -1))
+    u, s, v = da.linalg.svd_compressed(x, k=n_modes, compute=False)
+
+    u_rot = rotate(u)
+    dsk = dict(u_rot.dask)
+    dependencies, dependents = get_deps(dsk)
+    connected_roots, max_dependents = _connecting_to_roots(dependencies, dependents)
+
+    hashes = defaultdict(set)
+
+    for v in connected_roots.values():
+        hashes[tuple(sorted(v))].add(id(v))
+
+    assert set(map(len, hashes.values())) == {1}
+
+
 def test_connecting_to_roots_asym():
     dsk = {
         "a0": (f, 1),