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

Code before:
        ({"b": "mean"}, {"b": custom_mean}, False),
        ({"b": "sum"}, {"b": custom_sum}, True),
        (["mean", "sum"], [custom_mean, custom_sum], False),
        ({"b": ["mean", "sum"]}, {"b": [custom_mean, custom_sum]}, False),
    ],
)
def test_dataframe_groupby_agg_custom_sum(pandas_spec, dask_spec, check_dtype):
    df = pd.DataFrame({"g": [0, 0, 1] * 3, "b": [1, 2, 3] * 3})
    ddf = dd.from_pandas(df, npartitions=2)

    expected = df.groupby("g").aggregate(pandas_spec)
    result = ddf.groupby("g").aggregate(dask_spec)

    assert_eq(result, expected, check_dtype=check_dtype)


@pytest.mark.skipif(dd._dask_expr_enabled(), reason="Aggregation not supported")
@pytest.mark.parametrize(
    "pandas_spec, dask_spec",
    [
        ("mean", custom_mean),
        (["mean"], [custom_mean]),
        (["mean", "sum"], [custom_mean, custom_sum]),
    ],
)
def test_series_groupby_agg_custom_mean(pandas_spec, dask_spec):
    d = pd.DataFrame({"g": [0, 0, 1] * 3, "b": [1, 2, 3] * 3})
    a = dd.from_pandas(d, npartitions=2)

    expected = d["b"].groupby(d["g"]).aggregate(pandas_spec)
    result = a["b"].groupby(a["g"]).aggregate(dask_spec)

    assert_eq(result, expected, check_dtype=False)


@pytest.mark.skipif(dd._dask_expr_enabled(), reason="Aggregation not supported")
def test_groupby_agg_custom__name_clash_with_internal_same_column():
    """for a single input column only unique names are allowed"""
    d = pd.DataFrame({"g": [0, 0, 1] * 3, "b": [1, 2, 3] * 3})
    a = dd.from_pandas(d, npartitions=2)

    agg_func = dd.Aggregation("sum", lambda s: s.sum(), lambda s0: s0.sum())

    with pytest.raises(ValueError):
        a.groupby("g").aggregate({"b": [agg_func, "sum"]})


@pytest.mark.skipif(dd._dask_expr_enabled(), reason="Aggregation not supported")
def test_groupby_agg_custom__name_clash_with_internal_different_column():
    """custom aggregation functions can share the name of a builtin function"""
    d = pd.DataFrame({"g": [0, 0, 1] * 3, "b": [1, 2, 3] * 3, "c": [4, 5, 6] * 3})

Review comment (suggestion):
```suggestion
if not dd._dask_expr_enabled():
```

Code after:
    # Check that partitioning is preserved
    assert ddf.divisions == result.divisions

    # Check that no shuffling occurred.
    # The groupby operation should add only 1 task per partition
    assert len(result.dask) == (len(ddf.dask) + ddf.npartitions)

    expected = df.groupby(group_args).apply(apply_func, axis=0)

    # Compute on dask DataFrame without divisions (requires shuffling)
    result = ddf_no_divs.groupby(group_args).apply(apply_func, axis=0, meta=expected)

    assert_eq(expected, result, check_divisions=False)

    # Check that divisions were preserved (all None in this case)
    assert ddf_no_divs.divisions == result.divisions

    # Crude check to see if shuffling was performed.
    # The groupby operation should add only more than 1 task per partition
    assert len(result.dask) > (len(ddf_no_divs.dask) + ddf_no_divs.npartitions)


if dd._dask_expr_enabled():
    custom_mean = None
    custom_sum = None
else:
    custom_mean = dd.Aggregation(
        "mean",
        lambda s: (s.count(), s.sum()),
        lambda s0, s1: (s0.sum(), s1.sum()),
        lambda s0, s1: s1 / s0,
    )

    custom_sum = dd.Aggregation("sum", lambda s: s.sum(), lambda s0: s0.sum())


@pytest.mark.skipif(dd._dask_expr_enabled(), reason="Aggregation not supported")
@pytest.mark.parametrize(
    "pandas_spec, dask_spec, check_dtype",
    [
        ({"b": "mean"}, {"b": custom_mean}, False),
        ({"b": "sum"}, {"b": custom_sum}, True),
        (["mean", "sum"], [custom_mean, custom_sum], False),
        ({"b": ["mean", "sum"]}, {"b": [custom_mean, custom_sum]}, False),
    ],
)
def test_dataframe_groupby_agg_custom_sum(pandas_spec, dask_spec, check_dtype):
    df = pd.DataFrame({"g": [0, 0, 1] * 3, "b": [1, 2, 3] * 3})
    ddf = dd.from_pandas(df, npartitions=2)

    expected = df.groupby("g").aggregate(pandas_spec)