Repository: dask/dask
Pull request #11484 (2024-11-04 10:38:33)

Make sure that block_id and block_info don't create extra tasks
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


Putting the arguments into a blockwise layer will ensure that blockwise fusion kicks in. Currently, this creates one extra task per chunk, which is not great and generally interrupts blockwise fusion of different layers. map_overlap is using the block_info keyword internally, so this has immediate benefit, not only in 3rd party code

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 43m 28s ⏱️ + 2m 16s
 13 256 tests + 1   12 191 ✅ + 1   1 062 💤 ±0  3 ❌ ±0 
138 307 runs  +11  119 317 ✅ +13  18 987 💤  - 2  3 ❌ ±0 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/32470987162).

Results for commit c6a3d8af. ± Comparison against base commit 72f3ebb9.

[test-results]:data:application/gzip;base64,H4sIAEytKGcC/03Myw6DIBCF4VcxrLtgGItjX6ZBwITUS4OwavruHS9Fl+c7yf8RfRj8Ih4V4K0SSw7pXC5Hk8I8bbuWxMRv2n911//9XLK1Kypo4cRXeK8otSrWmzCwYQEf4xxZJEvM054mlM2xSxpahBOPNLVU7JLe9rVs53EMiYew2qAj0xMZ35GxioC7Gn3tatdplOiRoFHi+wNbgqxEGgEAAA==


Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index bb2e78e57..9ca312215 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -887,24 +887,38 @@ def map_blocks(
     extra_names = []
     # If func has block_id as an argument, construct an array of block IDs and
     # prepare to inject it.
+
+    def _getter_item(a, b, **kwargs):
+        return getter(a, b, **kwargs).item()
+
     if has_keyword(func, "block_id"):
+        # put block_id into a Blockwise layer so that we can fuse it
+        # with the other blockwise layers
+        block_id_arr = np.empty(tuple([len(c) for c in out.chunks]), dtype=np.object_)
+        for block_id in product(*(range(len(c)) for c in out.chunks)):
+            block_id_arr[block_id] = block_id
+
         block_id_name = "block-id-" + out.name
-        block_id_dsk = {
-            (block_id_name,) + block_id: block_id
-            for block_id in product(*(range(len(c)) for c in out.chunks))
-        }
-        block_id_array = Array(
-            block_id_dsk,
+        cs = tuple((1,) * len(c) for c in out.chunks)
+
+        dsk = graph_from_arraylike(
+            block_id_arr,
+            cs,
+            block_id_arr.shape,
             block_id_name,
-            chunks=tuple((1,) * len(c) for c in out.chunks),
-            dtype=np.object_,
+            getitem=_getter_item,
+            dtype=block_id_arr.dtype,
+            inline_array=True,
         )
+        block_id_array = Array(dsk, block_id_name, chunks=cs, dtype=np.object_)
         extra_argpairs.append((block_id_array, out_ind))
         extra_names.append("block_id")
 
     # If func has block_info as an argument, construct an array of block info
     # objects and prepare to inject it.
     if has_keyword(func, "block_info"):
+        # put block_info into a Blockwise layer so that we can fuse it
+        # with the other blockwise layers
         starts = {}
         num_chunks = {}
         shapes = {}
@@ -932,7 +946,7 @@ def map_blocks(
         out_starts = [cached_cumsum(c, initial_zero=True) for c in out.chunks]
 
         block_info_name = "block-info-" + out.name
-        block_info_dsk = {}
+        block_info_arr = np.empty(tuple([len(c) for c in out.chunks]), dtype=np.object_)
         for block_id in product(*(range(len(c)) for c in out.chunks)):
             # Get position of chunk, indexed by axis labels
             location = {out_ind[i]: loc for i, loc in enumerate(block_id)}
@@ -969,14 +983,19 @@ def map_blocks(
                 ),
                 "dtype": dtype,
             }
-            block_info_dsk[(block_info_name,) + block_id] = info
+            block_info_arr[block_id] = info
 
-        block_info = Array(
-            block_info_dsk,
+        cs = tuple((1,) * len(c) for c in out.chunks)
+        dsk = graph_from_arraylike(
+            block_info_arr,
+            cs,
+            block_info_arr.shape,
             block_info_name,
-            chunks=tuple((1,) * len(c) for c in out.chunks),
-            dtype=np.object_,
+            getitem=_getter_item,
+            dtype=block_info_arr.dtype,
+            inline_array=True,
         )
+        block_info = Array(dsk, block_info_name, chunks=cs, dtype=np.object_)
         extra_argpairs.append((block_info, out_ind))
         extra_names.append("block_info")
 
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
index 1bfe75590..73320fefa 100644
--- a/dask/array/overlap.py
+++ b/dask/array/overlap.py
@@ -9,7 +9,7 @@ from tlz.curried import map
 
 from dask.array import chunk
 from dask.array.core import Array, concatenate, map_blocks, unify_chunks
-from dask.array.creation import empty_like, full_like
+from dask.array.creation import empty_like, full_like, repeat
 from dask.array.numpy_compat import normalize_axis_tuple
 from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
@@ -243,8 +243,8 @@ def nearest(x, axis, depth):
         + (slice(None, None, None),) * (x.ndim - axis - 1)
     )
 
-    l = concatenate([x[left]] * depth, axis=axis)
-    r = concatenate([x[right]] * depth, axis=axis)
+    l = repeat(x[left], depth, axis=axis)
+    r = repeat(x[right], depth, axis=axis)
 
     l, r = _remove_overlap_boundaries(l, r, axis, depth)
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 5d99f2b01..6aaec6043 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1646,8 +1646,7 @@ def test_map_blocks_optimize_blockwise(func):
     optimized = optimize_blockwise(dsk)
 
     # Everything should be fused into a single layer.
-    # If the lambda includes block_info, there will be two layers.
-    assert len(optimized.layers) == len(dsk.layers) - 6
+    assert len(optimized.layers) == 1
 
 
 def test_repr():
diff --git a/dask/array/tests/test_map_blocks.py b/dask/array/tests/test_map_blocks.py
new file mode 100644
index 000000000..e27c74daf
--- /dev/null
+++ b/dask/array/tests/test_map_blocks.py
@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+import dask.array as da
+from dask.base import collections_to_dsk
+
+
+def test_map_blocks_block_id_fusion():
+    arr = da.ones((20, 10), chunks=(2, 5))
+
+    def dummy(x, block_id=None, block_info=None):
+        return x
+
+    result = arr.map_blocks(dummy).astype("f8")
+    dsk = collections_to_dsk([result])
+    assert len(dsk) == 20