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

Code before:
    else:
        assert d.groupby("a").by == "a"
        assert d.groupby(d["a"]).by == "a"
        assert d.groupby(d["a"] > 2).by._name == (d["a"] > 2)._name
    assert d.groupby(["a", "b"]).by == ["a", "b"]

    assert d.groupby([d["a"], d["b"]]).by == ["a", "b"]
    assert d.groupby([d["a"], "b"]).by == ["a", "b"]


def test_aggregate__single_element_groups(agg_func):
    spec = agg_func

    # nunique/cov is not supported in specs
    if spec in ("nunique", "cov", "corr"):
        return

    pdf = pd.DataFrame(
        {"a": [1, 1, 3, 3], "b": [4, 4, 16, 16], "c": [1, 1, 4, 4], "d": [1, 1, 3, 3]},
        columns=["c", "b", "a", "d"],
    )
    ddf = dd.from_pandas(pdf, npartitions=3)

    expected = pdf.groupby(["a", "d"]).agg(spec)

    # NOTE: for std the result is not recast of the original dtype
    if spec in {"mean", "var"}:
        expected = expected.astype(float)

    shuffle_method = (
        {"shuffle_method": "tasks", "split_out": 2} if agg_func == "median" else {}
    )
    assert_eq(expected, ddf.groupby(["a", "d"]).agg(spec, **shuffle_method))


def test_aggregate_build_agg_args__reuse_of_intermediates():
    """Aggregate reuses intermediates. For example, with sum, count, and mean
    the sums and counts are only calculated once across the graph and reused to
    compute the mean.
    """
    from dask.dataframe.groupby import _build_agg_args

    no_mean_spec = [("foo", "sum", "input"), ("bar", "count", "input")]

    with_mean_spec = [
        ("foo", "sum", "input"),
        ("bar", "count", "input"),
        ("baz", "mean", "input"),
    ]

    no_mean_chunks, no_mean_aggs, no_mean_finalizers = _build_agg_args(no_mean_spec)

Review comment (suggestion):
```suggestion
    # NOTE: for std the result is not recast to the original dtype
```

Code after:
    else:
        assert d.groupby("a").by == "a"
        assert d.groupby(d["a"]).by == "a"
        assert d.groupby(d["a"] > 2).by._name == (d["a"] > 2)._name
    assert d.groupby(["a", "b"]).by == ["a", "b"]

    assert d.groupby([d["a"], d["b"]]).by == ["a", "b"]
    assert d.groupby([d["a"], "b"]).by == ["a", "b"]


def test_aggregate__single_element_groups(agg_func):
    spec = agg_func

    # nunique/cov is not supported in specs
    if spec in ("nunique", "cov", "corr"):
        return

    pdf = pd.DataFrame(
        {"a": [1, 1, 3, 3], "b": [4, 4, 16, 16], "c": [1, 1, 4, 4], "d": [1, 1, 3, 3]},
        columns=["c", "b", "a", "d"],
    )
    ddf = dd.from_pandas(pdf, npartitions=3)

    expected = pdf.groupby(["a", "d"]).agg(spec)

    # NOTE: for std the result is not recast to the original dtype
    if spec in {"mean", "var"}:
        expected = expected.astype(float)

    shuffle_method = (
        {"shuffle_method": "tasks", "split_out": 2} if agg_func == "median" else {}
    )
    assert_eq(expected, ddf.groupby(["a", "d"]).agg(spec, **shuffle_method))


def test_aggregate_build_agg_args__reuse_of_intermediates():
    """Aggregate reuses intermediates. For example, with sum, count, and mean
    the sums and counts are only calculated once across the graph and reused to
    compute the mean.
    """
    from dask.dataframe.groupby import _build_agg_args

    no_mean_spec = [("foo", "sum", "input"), ("bar", "count", "input")]

    with_mean_spec = [
        ("foo", "sum", "input"),
        ("bar", "count", "input"),
        ("baz", "mean", "input"),
    ]

    no_mean_chunks, no_mean_aggs, no_mean_finalizers = _build_agg_args(no_mean_spec)