Function test_map_blocks_da_ds_with_template

xarray/tests/test_dask.py:1410–1427

(obj)

@pytest.mark.parametrize("obj", [make_da(), make_ds()])
def test_map_blocks_da_ds_with_template(obj):
    func = lambda x: x.isel(x=[1])
    # a simple .isel(x=[1, 5, 9]) puts all those in a single chunk.
    template = xr.concat([obj.isel(x=[i]) for i in [1, 5, 9]], data_vars=None, dim="x")
    with raise_if_dask_computes():
        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 {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)


def test_map_blocks_roundtrip_string_index():