Repository: pydata/xarray

Issue:
Task graphs on `.map_blocks` with many chunks can be huge
### What happened?

I'm getting task graphs > 1GB, I think possibly because the full indexes are being included in every task?

### What did you expect to happen?

Only the relevant sections of the index would be included

### Minimal Complete Verifiable Example

```Python
da = xr.tutorial.load_dataset('air_temperature')

# Dropping the index doesn't generally matter that much...

len(cloudpickle.dumps(da.chunk(lat=1, lon=1)))
# 15569320

len(cloudpickle.dumps(da.chunk().drop_vars(da.indexes)))
# 15477313

# But with `.map_blocks`, it really matters — it's really big with the indexes, and the same size without:


len(cloudpickle.dumps(da.chunk(lat=1, lon=1).map_blocks(lambda x: x)))
# 79307120

len(cloudpickle.dumps(da.chunk(lat=1, lon=1).drop_vars(da.indexes).map_blocks(lambda x: x)))
# 16016173
```


### MVCE confirmation

- [X] Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
- [X] Complete example — the example is self-contained, including all data and the text of any traceback.
- [X] Verifiable example — the example copy & pastes into an IPython prompt or [Binder notebook](https://mybinder.org/v2/gh/pydata/xarray/main?urlpath=lab/tree/doc/examples/blank_template.ipynb), returning the result.
- [X] New issue — a search of GitHub Issues suggests this is not a duplicate.
- [X] Recent environment — the issue occurs with the latest version of xarray and its dependencies.

### Relevant log output

_No response_

### Anything else we need to know?

_No response_

### Environment

<details>


INSTALLED VERSIONS
------------------
commit: None
python: 3.9.18 (main, Aug 24 2023, 21:19:58)
[Clang 14.0.3 (clang-1403.0.22.14.1)]
python-bits: 64
OS: Darwin
OS-release: 22.6.0
machine: arm64
processor: arm
byteorder: little
LC_ALL: en_US.UTF-8
LANG: None
LOCALE: ('en_US', 'UTF-8')
libhdf5: 1.12.2
libnetcdf: None

xarray: 2023.10.1
pandas: 2.1.1
numpy: 1.26.1
scipy: 1.11.1
netCDF4: None
pydap: None
h5netcdf: 1.1.0
h5py: 3.8.0
Nio: None
zarr: 2.16.0
cftime: 1.6.2
nc_time_axis: None
PseudoNetCDF: None
iris: None
bottleneck: 1.3.7
dask: 2023.5.0
distributed: 2023.5.0
matplotlib: 3.6.0
cartopy: None
seaborn: 0.12.2
numbagg: 0.6.0
fsspec: 2022.8.2
cupy: None
pint: 0.22
sparse: 0.14.0
flox: 0.7.2
numpy_groupies: 0.9.22
setuptools: 68.1.2
pip: 23.2.1
conda: None
pytest: 7.4.0
mypy: 1.6.1
IPython: 8.14.0
sphinx: 5.2.1



</details>


Discussion:
dcherian: > Is this something that you all would consider merging?

Yes. This would extend what was done in https://github.com/pydata/xarray/pull/8412.

I *think* I initially chose this approach to make the graph absolutely embarrassingly parallel.

Can you add a test or confirm that dask has a test that makes sure the ordering scheduling is "nice" for this kind of case (widely duplicated dependency)? 
phofl: > Can you add a test to dask or confirm that dask has a test that makes sure the ordering/scheduling is "nice" for this kind of case (widely duplicated dependency)?

This is the same pattern as we have with the open_dataset tasks and in the new dask array take implementation, we have existing tests covering these 2 patterns already
phofl: > Ya this pattern used to suck, so I avoided it. But good to know it's fixed. Thanks for the PR!

TLDR is: If your task does not have a callable, it will be ignored during ordering

Patch:
diff --git a/xarray/core/parallel.py b/xarray/core/parallel.py
index 84728007..a0dfe568 100644
--- a/xarray/core/parallel.py
+++ b/xarray/core/parallel.py
@@ -25,7 +25,6 @@ class ExpectedDict(TypedDict):
     shapes: dict[Hashable, int]
     coords: set[Hashable]
     data_vars: set[Hashable]
-    indexes: dict[Hashable, Index]
 
 
 def unzip(iterable):
@@ -337,6 +336,7 @@ def map_blocks(
         kwargs: dict,
         arg_is_array: Iterable[bool],
         expected: ExpectedDict,
+        expected_indexes: dict[Hashable, Index],
     ):
         """
         Wrapper function that receives datasets in args; converts to dataarrays when necessary;
@@ -372,7 +372,7 @@ def map_blocks(
 
             # ChainMap wants MutableMapping, but xindexes is Mapping
             merged_indexes = collections.ChainMap(
-                expected["indexes"],
+                expected_indexes,
                 merged_coordinates.xindexes,  # type: ignore[arg-type]
             )
             expected_index = merged_indexes.get(name, None)
@@ -412,6 +412,7 @@ def map_blocks(
     try:
         import dask
         import dask.array
+        from dask.base import tokenize
         from dask.highlevelgraph import HighLevelGraph
 
     except ImportError:
@@ -551,6 +552,20 @@ def map_blocks(
             for isxr, arg in zip(is_xarray, npargs, strict=True)
         ]
 
+        # only include new or modified indexes to minimize duplication of data
+        indexes = {
+            dim: coordinates.xindexes[dim][
+                _get_chunk_slicer(dim, chunk_index, output_chunk_bounds)
+            ]
+            for dim in (new_indexes | modified_indexes)
+        }
+
+        tokenized_indexes: dict[Hashable, str] = {}
+        for k, v in indexes.items():
+            tokenized_v = tokenize(v)
+            graph[f"{k}-coordinate-{tokenized_v}"] = v
+            tokenized_indexes[k] = f"{k}-coordinate-{tokenized_v}"
+
         # raise nice error messages in _wrapper
         expected: ExpectedDict = {
             # input chunk 0 along a dimension maps to output chunk 0 along the same dimension
@@ -562,17 +577,18 @@ def map_blocks(
             },
             "data_vars": set(template.data_vars.keys()),
             "coords": set(template.coords.keys()),
-            # only include new or modified indexes to minimize duplication of data, and graph size.
-            "indexes": {
-                dim: coordinates.xindexes[dim][
-                    _get_chunk_slicer(dim, chunk_index, output_chunk_bounds)
-                ]
-                for dim in (new_indexes | modified_indexes)
-            },
         }
 
         from_wrapper = (gname,) + chunk_tuple
-        graph[from_wrapper] = (_wrapper, func, blocked_args, kwargs, is_array, expected)
+        graph[from_wrapper] = (
+            _wrapper,
+            func,
+            blocked_args,
+            kwargs,
+            is_array,
+            expected,
+            (dict, [[k, v] for k, v in tokenized_indexes.items()]),
+        )
 
         # mapping from variable name to dask graph key
         var_key_map: dict[Hashable, str] = {}


Test patch:
diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py
index a46a9d43..cc795b75 100644
--- a/xarray/tests/test_dask.py
+++ b/xarray/tests/test_dask.py
@@ -14,6 +14,7 @@ import xarray as xr
 from xarray import DataArray, Dataset, Variable
 from xarray.core import duck_array_ops
 from xarray.core.duck_array_ops import lazy_array_equiv
+from xarray.core.indexes import PandasIndex
 from xarray.testing import assert_chunks_equal
 from xarray.tests import (
     assert_allclose,
@@ -1375,6 +1376,13 @@ def test_map_blocks_da_ds_with_template(obj):
         actual = xr.map_blocks(func, obj, template=template)
     assert_identical(actual, template)
 
+    # Check that indexes are written into the graph directly
+    dsk = dict(actual.__dask_graph__())
+    assert len({k for k in dsk if "x-coordinate" in k})
+    assert all(
+        isinstance(v, PandasIndex) for k, v in dsk.items() if "x-coordinate" in k
+    )
+
     with raise_if_dask_computes():
         actual = obj.map_blocks(func, template=template)
     assert_identical(actual, template)