Function test_array_vs_dataframe

dask/tests/test_order.py:1395–1428

(optimize)

    ],
)
def test_array_vs_dataframe(optimize):
    xr = pytest.importorskip("xarray")
    pytest.importorskip("dask.dataframe")

    import dask.array as da

    size = 5000
    ds = xr.Dataset(
        dict(
            anom_u=(
                ["time", "face", "j", "i"],
                da.random.random((size, 1, 987, 1920), chunks=(10, 1, -1, -1)),
            ),
            anom_v=(
                ["time", "face", "j", "i"],
                da.random.random((size, 1, 987, 1920), chunks=(10, 1, -1, -1)),
            ),
        )
    )

    quad = ds**2
    quad["uv"] = ds.anom_u * ds.anom_v
    mean = quad.mean("time")
    diag_array = diagnostics(
        collections_to_expr([mean], optimize_graph=optimize).optimize().__dask_graph__()
    )
    diag_df = diagnostics(
        collections_to_expr([mean.to_dask_dataframe()], optimize_graph=optimize)
        .optimize()
        .__dask_graph__()
    )
    assert max(diag_df[1]) == 15
    assert max(diag_array[1]) == 38
    assert max(diag_array[1]) < 50


def test_anom_mean():