Function test_map_overlap_multiple_dataframes

dask/dataframe/tests/test_rolling.py:96–153

(
    use_dask_input,
    npartitions,
    enforce_metadata,
    transform_divisions,
    align_dataframes,
    overlap_setup,
)

    ],
)
def test_map_overlap_multiple_dataframes(
    use_dask_input,
    npartitions,
    enforce_metadata,
    transform_divisions,
    align_dataframes,
    overlap_setup,
):
    dataframe, before, after = overlap_setup

    ddf = dataframe
    ddf2 = dataframe * 2
    if use_dask_input:
        ddf = dd.from_pandas(ddf, npartitions)
        ddf2 = dd.from_pandas(ddf2, 2 if align_dataframes else npartitions)

    def get_shifted_sum_arg(overlap):
        return (
            overlap.seconds - 1 if isinstance(overlap, datetime.timedelta) else overlap
        )

    before_shifted_sum, after_shifted_sum = get_shifted_sum_arg(
        before
    ), get_shifted_sum_arg(after)

    # DataFrame
    res = dd.map_overlap(
        shifted_sum,
        ddf,
        before,
        after,
        before_shifted_sum,
        after_shifted_sum,
        ddf2,
        align_dataframes=align_dataframes,
        transform_divisions=transform_divisions,
        enforce_metadata=enforce_metadata,
    )
    sol = shifted_sum(dataframe, before_shifted_sum, after_shifted_sum, dataframe * 2)
    assert_eq(res, sol)

    # Series
    res = dd.map_overlap(
        shifted_sum,
        ddf.b,
        before,
        after,
        before_shifted_sum,
        after_shifted_sum,
        ddf2.b,
        align_dataframes=align_dataframes,
        transform_divisions=transform_divisions,
        enforce_metadata=enforce_metadata,
    )
    sol = shifted_sum(
        dataframe.b, before_shifted_sum, after_shifted_sum, dataframe.b * 2
    )
    assert_eq(res, sol)