Repository: dask/dask
Pull request #11478 (2024-10-31 16:40:30)

Remove dead code from make_blockwise
From all I can tell, this is dead code. The function `make_blockwise_graph` is used in exactly one place and the kwargs `deserialize`, `return_key_deps` and `func_future_args` are never used. If we remove them, this other code is also dead

fjetter: I did check that this isn't used in xarray but haven't checked in other projects
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 47m 23s ⏱️ + 2m 50s
 13 254 tests  -  1   12 192 ✅  -  1   1 062 💤 ±0  0 ❌ ±0 
138 285 runs   - 13  119 295 ✅  - 11  18 990 💤  - 2  0 ❌ ±0 

Results for commit 97f22612. ± Comparison against base commit 0cb1cb1f.

<details>
  <summary>This pull request <b>removes</b> 1 test.</summary>

```
dask.tests.test_utils ‑ test_stringify_collection_keys
```
</details>

[test-results]:data:application/gzip;base64,H4sIAKe7I2cC/02Myw6DIBQFf8Ww7gIugtz+TGNRElIfDY+V6b8XFMXlzMmZjRg7jZ48G8YfDfHRhkpDdH2w67KzbLNKazh2EO3JLx+1zhIYQpUf+82SyupMb6fk6CVG51ZXjIvLkVagROErncpYZUkrxPN4T+98L+t1nm1IQLAzAJKBMJxp2tGB0VYapFIYwVGjVqzHQb3J7w/kUw8mGgEAAA==


Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index eb6f40723..bb2e78e57 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -3223,6 +3223,12 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
     normalize_chunks: for full docstring and parameters
     """
     if previous_chunks is not None:
+        # rioxarray is passing ((1, ), (x,)) for shapes like (100, 5x),
+        # so add this compat code for now
+        # https://github.com/corteva/rioxarray/pull/820
+        previous_chunks = (
+            c[0] if isinstance(c, tuple) and len(c) == 1 else c for c in previous_chunks
+        )
         previous_chunks = _convert_int_chunk_to_tuple(shape, previous_chunks)
     chunks = list(chunks)
 
diff --git a/dask/blockwise.py b/dask/blockwise.py
index 1acd4d6ec..141eb3db0 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -15,13 +15,7 @@ from dask.core import flatten, ishashable, keys_in_tasks, reverse_dict
 from dask.highlevelgraph import HighLevelGraph, Layer
 from dask.optimization import SubgraphCallable, fuse
 from dask.typing import Graph, Key
-from dask.utils import (
-    _deprecated,
-    apply,
-    ensure_dict,
-    homogeneous_deepmap,
-    stringify_collection_keys,
-)
+from dask.utils import _deprecated, apply, ensure_dict, homogeneous_deepmap
 
 
 class BlockwiseDep:
@@ -768,9 +762,6 @@ def make_blockwise_graph(
     new_axes=None,
     output_blocks=None,
     dims=None,
-    deserializing=False,
-    func_future_args=None,
-    return_key_deps=False,
     io_deps=None,
 ):
     """Tensor operation
@@ -883,12 +874,6 @@ def make_blockwise_graph(
     io_deps = io_deps or {}
     argpairs = list(toolz.partition(2, arrind_pairs))
 
-    if return_key_deps:
-        key_deps = {}
-
-    if deserializing:
-        from distributed.protocol.serialize import to_serialize
-
     if concatenate is True:
         from dask.array.core import concatenate_axes as concatenate
 
@@ -924,10 +909,7 @@ def make_blockwise_graph(
         args = []
         for cmap, axes, (arg, ind) in zip(coord_maps, concat_axes, argpairs):
             if ind is None:
-                if deserializing:
-                    args.append(stringify_collection_keys(arg))
-                else:
-                    args.append(arg)
+                args.append(arg)
             else:
                 arg_coords = tuple(coords[c] for c in cmap)
                 if axes:
@@ -947,44 +929,14 @@ def make_blockwise_graph(
                     # we are replacing here
                     idx = tups[1:]
                     args.append(io_deps[arg].get(idx, idx))
-                elif deserializing:
-                    args.append(stringify_collection_keys(tups))
                 else:
                     args.append(tups)
         out_key = (output,) + out_coords
 
-        if deserializing:
-            deps.update(func_future_args)
-            args += list(func_future_args)
-
-        # Construct a function/args/kwargs dict if we
-        # do not have a nested task (i.e. concatenate=False).
-        # TODO: Avoid using the iterate_collection-version
-        # of to_serialize if we know that are no embedded
-        # Serialized/Serialize objects in args and/or kwargs.
-        if deserializing and isinstance(func, bytes):
-            dsk[out_key] = {"function": func, "args": to_serialize(args)}
-        else:
-            args.insert(0, func)
-            val = tuple(args)
-            # May still need to serialize (if concatenate=True)
-            dsk[out_key] = to_serialize(val) if deserializing else val
-
-        if return_key_deps:
-            key_deps[out_key] = deps
+        args.insert(0, func)
+        dsk[out_key] = tuple(args)
 
-    if return_key_deps:
-        # Add valid-key dependencies from io_deps
-        for key, io_dep in io_deps.items():
-            if io_dep.produces_keys:
-                for out_coords in output_blocks:
-                    key = (output,) + out_coords
-                    valid_key_dep = io_dep[out_coords]
-                    key_deps[key] |= {valid_key_dep}
-
-        return dsk, key_deps
-    else:
-        return dsk
+    return dsk
 
 
 def lol_product(head, values):
diff --git a/dask/utils.py b/dask/utils.py
index fc9322947..00e4f09ca 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -2064,29 +2064,6 @@ def stringify(obj, exclusive: Iterable | None = None):
     return obj
 
 
-def stringify_collection_keys(obj):
-    """Convert all collection keys in ``obj`` to strings.
-
-    This is a specialized version of ``stringify()`` that only converts keys
-    of the form: ``("a string", ...)``
-    """
-
-    typ = type(obj)
-    if typ is tuple and obj:
-        obj0 = obj[0]
-        if type(obj0) is str or type(obj0) is bytes:
-            return stringify(obj)
-        if callable(obj0):
-            return (obj0,) + tuple(stringify_collection_keys(x) for x in obj[1:])
-    if typ is list:
-        return [stringify_collection_keys(v) for v in obj]
-    if typ is dict:
-        return {k: stringify_collection_keys(v) for k, v in obj.items()}
-    if typ is tuple:  # If the tuple itself isn't a key, check its elements
-        return tuple(stringify_collection_keys(v) for v in obj)
-    return obj
-
-
 class cached_property(functools.cached_property):
     """Read only version of functools.cached_property."""
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 121f8a698..5d99f2b01 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -54,9 +54,11 @@ from dask.array.numpy_compat import NUMPY_GE_200
 from dask.array.reshape import _not_implemented_message
 from dask.array.utils import assert_eq, same_keys
 from dask.base import compute_as_if_collection, tokenize
-from dask.blockwise import broadcast_dimensions
-from dask.blockwise import make_blockwise_graph as top
-from dask.blockwise import optimize_blockwise
+from dask.blockwise import (
+    broadcast_dimensions,
+    make_blockwise_graph,
+    optimize_blockwise,
+)
 from dask.delayed import Delayed, delayed
 from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
 from dask.layers import Blockwise
@@ -90,14 +92,14 @@ def test_graph_from_arraylike(inline_array):
 
 
 def test_top():
-    assert top(inc, "z", "ij", "x", "ij", numblocks={"x": (2, 2)}) == {
+    assert make_blockwise_graph(inc, "z", "ij", "x", "ij", numblocks={"x": (2, 2)}) == {
         ("z", 0, 0): (inc, ("x", 0, 0)),
         ("z", 0, 1): (inc, ("x", 0, 1)),
         ("z", 1, 0): (inc, ("x", 1, 0)),
         ("z", 1, 1): (inc, ("x", 1, 1)),
     }
 
-    assert top(
+    assert make_blockwise_graph(
         add, "z", "ij", "x", "ij", "y", "ij", numblocks={"x": (2, 2), "y": (2, 2)}
     ) == {
         ("z", 0, 0): (add, ("x", 0, 0), ("y", 0, 0)),
@@ -106,7 +108,7 @@ def test_top():
         ("z", 1, 1): (add, ("x", 1, 1), ("y", 1, 1)),
     }
 
-    assert top(
+    assert make_blockwise_graph(
         dotmany, "z", "ik", "x", "ij", "y", "jk", numblocks={"x": (2, 2), "y": (2, 2)}
     ) == {
         ("z", 0, 0): (dotmany, [("x", 0, 0), ("x", 0, 1)], [("y", 0, 0), ("y", 1, 0)]),
@@ -115,13 +117,13 @@ def test_top():
         ("z", 1, 1): (dotmany, [("x", 1, 0), ("x", 1, 1)], [("y", 0, 1), ("y", 1, 1)]),
     }
 
-    assert top(identity, "z", "", "x", "ij", numblocks={"x": (2, 2)}) == {
-        ("z",): (identity, [[("x", 0, 0), ("x", 0, 1)], [("x", 1, 0), ("x", 1, 1)]])
-    }
+    assert make_blockwise_graph(
+        identity, "z", "", "x", "ij", numblocks={"x": (2, 2)}
+    ) == {("z",): (identity, [[("x", 0, 0), ("x", 0, 1)], [("x", 1, 0), ("x", 1, 1)]])}
 
 
 def test_top_supports_broadcasting_rules():
-    assert top(
+    assert make_blockwise_graph(
         add, "z", "ij", "x", "ij", "y", "ij", numblocks={"x": (1, 2), "y": (2, 1)}
     ) == {
         ("z", 0, 0): (add, ("x", 0, 0), ("y", 0, 0)),
@@ -132,7 +134,9 @@ def test_top_supports_broadcasting_rules():
 
 
 def test_top_literals():
-    assert top(add, "z", "ij", "x", "ij", 123, None, numblocks={"x": (2, 2)}) == {
+    assert make_blockwise_graph(
+        add, "z", "ij", "x", "ij", 123, None, numblocks={"x": (2, 2)}
+    ) == {
         ("z", 0, 0): (add, ("x", 0, 0), 123),
         ("z", 0, 1): (add, ("x", 0, 1), 123),
         ("z", 1, 0): (add, ("x", 1, 0), 123),
@@ -216,7 +220,7 @@ def test_chunked_dot_product():
     getx = graph_from_arraylike(x, (5, 5), shape=(20, 20), name="x")
     geto = graph_from_arraylike(o, (5, 5), shape=(20, 20), name="o")
 
-    result = top(
+    result = make_blockwise_graph(
         dotmany, "out", "ik", "x", "ij", "o", "jk", numblocks={"x": (4, 4), "o": (4, 4)}
     )
 
@@ -232,7 +236,7 @@ def test_chunked_transpose_plus_one():
     getx = graph_from_arraylike(x, (5, 5), shape=(20, 20), name="x")
 
     f = lambda x: x.T + 1
-    comp = top(f, "out", "ij", "x", "ji", numblocks={"x": (4, 4)})
+    comp = make_blockwise_graph(f, "out", "ij", "x", "ji", numblocks={"x": (4, 4)})
 
     dsk = merge(getx, comp)
     out = dask.get(dsk, [[("out", i, j) for j in range(4)] for i in range(4)])
@@ -2913,6 +2917,19 @@ def test_normalize_chunks():
         normalize_chunks(((5,), (5,)), (5,))
 
 
+def test_single_element_tuple():
+    assert normalize_chunks(
+        (100, "auto"), (500, 500_000), dtype=np.int64, previous_chunks=((1,), (500,))
+    ) == (
+        (100,) * 5,
+        (
+            167_500,
+            167_500,
+            165_000,
+        ),
+    )
+
+
 def test_align_chunks_to_previous_chunks():
     chunks = normalize_chunks(
         "auto", shape=(2000,), previous_chunks=(512,), limit="600 B", dtype=np.uint8
diff --git a/dask/tests/test_utils.py b/dask/tests/test_utils.py
index 8aba2b3f3..5a69e415f 100644
--- a/dask/tests/test_utils.py
+++ b/dask/tests/test_utils.py
@@ -43,7 +43,6 @@ from dask.utils import (
     random_state_data,
     skip_doctest,
     stringify,
-    stringify_collection_keys,
     takes_multiple_arguments,
     tmpfile,
     typename,
@@ -783,17 +782,6 @@ def test_stringify():
     assert dsk[("y", 1)][1][0] == "('z', 1)"
 
 
-def test_stringify_collection_keys():
-    obj = "Hello"
-    assert stringify_collection_keys(obj) is obj
-
-    obj = [("a", 0), (b"a", 0), (1, 1)]
-    res = stringify_collection_keys(obj)
-    assert res[0] == str(obj[0])
-    assert res[1] == str(obj[1])
-    assert res[2] == obj[2]
-
-
 @pytest.mark.parametrize(
     "n,expect",
     [