From a4bda14fae54a33b6ee15095e27b5f89b6fc935f Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Mon, 18 Dec 2023 20:30:00 -0700 Subject: [PATCH] Add test --- xarray/tests/test_dask.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index 137d6020829..8043cfa81d6 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -1746,3 +1746,20 @@ def test_new_index_var_computes_once(): data = dask.array.from_array(np.array([100, 200])) with raise_if_dask_computes(max_computes=1): Dataset(coords={"z": ("z", data)}) + + +def test_minimize_graph_size(): + import cloudpickle + + # regression test for https://github.com/pydata/xarray/issues/8409 + ds = Dataset( + {"foo": (("x", "y", "z"), dask.array.ones((120, 120, 120)))}, + coords={"x": np.arange(120), "y": np.arange(120), "z": np.arange(120)}, + ).chunk(x=20, y=20, z=1) + + actual = len(cloudpickle.dumps(ds.map_blocks(lambda x: x))) + expected = len(cloudpickle.dumps(ds.drop_vars(ds.dims).map_blocks(lambda x: x))) + + # prior to https://github.com/pydata/xarray/pull/8412 + # actual is ~5x expected + assert actual < 2 * expected