Function test_minimize_graph_size

xarray/tests/test_dask.py:1832–1854

()

def test_minimize_graph_size():
    # regression test for https://github.com/pydata/xarray/issues/8409
    ds = Dataset(
        {
            "foo": (
                ("x", "y", "z"),
                dask.array.ones((120, 120, 120), chunks=(20, 20, 1)),
            )
        },
        coords={"x": np.arange(120), "y": np.arange(120), "z": np.arange(120)},
    )

    mapped = ds.map_blocks(lambda x: x)
    graph = dict(mapped.__dask_graph__())

    numchunks = {k: len(v) for k, v in ds.chunksizes.items()}
    for var in "xyz":
        actual = len([key for key in graph if var in key[0]])
        # assert that we only include each chunk of an index variable
        # is only included once, not the product of number of chunks of
        # all the other dimensions.
        # e.g. previously for 'x',  actual == numchunks['y'] * numchunks['z']
        assert actual == numchunks[var], (actual, numchunks[var])


def test_idxmin_chunking():