Function test_graph_manipulation

xarray/tests/test_dask.py:1793–1822

dask.graph_manipulation passes an optional parameter, "rename", to the rebuilder function returned by __dask_postperist__; also, the dsk passed to the rebuilder is a HighLevelGraph whereas with dask.persist() and dask.optimize() it's a plain dict.

()

def test_graph_manipulation():
    """dask.graph_manipulation passes an optional parameter, "rename", to the rebuilder
    function returned by __dask_postperist__; also, the dsk passed to the rebuilder is
    a HighLevelGraph whereas with dask.persist() and dask.optimize() it's a plain dict.
    """
    import dask.graph_manipulation as gm

    v = Variable(["x"], [1, 2]).chunk(-1).chunk(1) * 2
    da = DataArray(v)
    ds = Dataset({"d1": v[0], "d2": v[1], "d3": ("x", [3, 4])})

    v2, da2, ds2 = gm.clone(v, da, ds)

    assert_equal(v2, v)
    assert_equal(da2, da)
    assert_equal(ds2, ds)

    for a, b in ((v, v2), (da, da2), (ds, ds2)):
        assert a.__dask_layers__() != b.__dask_layers__()
        assert len(a.__dask_layers__()) == len(b.__dask_layers__())
        assert a.__dask_graph__().keys() != b.__dask_graph__().keys()  # type: ignore[union-attr]
        assert len(a.__dask_graph__()) == len(b.__dask_graph__())  # type: ignore[arg-type]
        assert a.__dask_graph__().layers.keys() != b.__dask_graph__().layers.keys()  # type: ignore[union-attr]
        assert len(a.__dask_graph__().layers) == len(b.__dask_graph__().layers)  # type: ignore[union-attr]

    # Above we performed a slice operation; adding the two slices back together creates
    # a diamond-shaped dependency graph, which in turn will trigger a collision in layer
    # names if we were to use HighLevelGraph.cull() instead of
    # HighLevelGraph.cull_layers() in Dataset.__dask_postpersist__().
    assert_equal(ds2.d1 + ds2.d2, ds.d1 + ds.d2)


def test_new_index_var_computes_once():