Repository: dask/dask
File: dask/dataframe/tests/test_groupby.py (Python)

Code before:
    if indexed:
        pdf = pdf.set_index("group")

    ddf = dd.from_pandas(pdf, npartitions)

    with pytest.warns(UserWarning):
        # DataFrame
        assert_eq(
            pdf.groupby("group").transform(lambda series: series - series.mean()),
            ddf.groupby("group").transform(lambda series: series - series.mean()),
        )

        # Series
        assert_eq(
            pdf.groupby("group")["value"].transform(
                lambda series: series - series.mean()
            ),
            ddf.groupby("group")["value"].transform(
                lambda series: series - series.mean()
            ),
        )


@pytest.mark.parametrize("dropna", [False, True, None])
def test_groupby_dropna_pandas(dropna):

    # The `dropna` arg is not currently supported by pandas
    # (See #https://github.com/pandas-dev/pandas/pull/21669)
    # Dask supports the argument for the cudf backend,
    # but should raise an error for pandas.

    # TODO: Expand test when `dropna`is supported in pandas.

    df = pd.DataFrame(
        {"a": [1, 2, 3, 4, None, None, 7, 8], "e": [4, 5, 6, 3, 2, 1, 0, 0]}
    )
    ddf = dd.from_pandas(df, npartitions=3)

    if dropna is None:
        dask_result = ddf.groupby("a").e.sum()
        pd_result = df.groupby("a").e.sum()
        assert_eq(dask_result, pd_result)
    else:
        with pytest.raises(TypeError):
            ddf.groupby("a", dropna=dropna)
        with pytest.raises(TypeError):
            df.groupby("a", dropna=dropna)


@pytest.mark.parametrize("dropna", [False, True, None])
@pytest.mark.parametrize("by", ["a", "c", "d", ["a", "b"], ["a", "c"], ["a", "d"]])

Review comment (bug):
I'd prefer that this test implements the behavior we expect pandas to adopt, and be marked with an `xfail`.

Then, when pandas does add support for `dropna` this test will start failing and we'll just have to remove the xfail (for new pandas).

So, can you remove the `dropna=None` case and change the TypeErrors to the expected result?

Code after:
    if indexed:
        pdf = pdf.set_index("group")

    ddf = dd.from_pandas(pdf, npartitions)

    with pytest.warns(UserWarning):
        # DataFrame
        assert_eq(
            pdf.groupby("group").transform(lambda series: series - series.mean()),
            ddf.groupby("group").transform(lambda series: series - series.mean()),
        )

        # Series
        assert_eq(
            pdf.groupby("group")["value"].transform(
                lambda series: series - series.mean()
            ),
            ddf.groupby("group")["value"].transform(
                lambda series: series - series.mean()
            ),
        )


@pytest.mark.xfail(reason="dropna kwarg not supported in pandas groupby.")
@pytest.mark.parametrize("dropna", [False, True])
def test_groupby_dropna_pandas(dropna):

    # The `dropna` arg is not currently supported by pandas
    # (See #https://github.com/pandas-dev/pandas/pull/21669)
    # Dask supports the argument for the cudf backend,
    # but passing it to the pandas backend will fail.

    # TODO: Expand test when `dropna` is supported in pandas.
    #       (See: `test_groupby_dropna_cudf`)

    df = pd.DataFrame(
        {"a": [1, 2, 3, 4, None, None, 7, 8], "e": [4, 5, 6, 3, 2, 1, 0, 0]}
    )
    ddf = dd.from_pandas(df, npartitions=3)

    dask_result = ddf.groupby("a", dropna=dropna)
    pd_result = df.groupby("a", dropna=dropna)
    assert_eq(dask_result, pd_result)


@pytest.mark.parametrize("dropna", [False, True, None])
@pytest.mark.parametrize("by", ["a", "c", "d", ["a", "b"], ["a", "c"], ["a", "d"]])
def test_groupby_dropna_cudf(dropna, by):

    # NOTE: This test requires cudf/dask_cudf, and will
    # be skipped by non-GPU CI