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

Code before:
    with_mean_chunks, with_mean_aggs, with_mean_finalizers = _build_agg_args(
        with_mean_spec
    )

    assert len(no_mean_chunks) == len(with_mean_chunks)
    assert len(no_mean_aggs) == len(with_mean_aggs)

    assert len(no_mean_finalizers) == len(no_mean_spec)
    assert len(with_mean_finalizers) == len(with_mean_spec)


def test_aggregate_dask():
    dask_holder = collections.namedtuple("dask_holder", ["dask"])
    get_agg_dask = lambda obj: dask_holder(
        {
            k: v
            for (k, v) in obj.dask.items()
            # Skip "chunk" tasks, because they include
            # SubgraphCallable object with non-deterministic
            # (uuid-based) function names
            if (k[0].startswith("aggregate") and "-chunk-" not in k[0])
        }
    )

    specs = [
        {"b": {"c": "mean"}, "c": {"a": "max", "b": "min"}},
        {"b": "mean", "c": ["min", "max"]},
        [
            "sum",
            "mean",
            "min",
            "max",
            "count",
            "size",
        ],
        [
            "std",
            "var",
            "first",
            "last",
            "prod",
        ],
        "sum",
        "mean",
        "min",
        "max",
        "count",
        "std",
        "var",
        "first",
        "last",

Review comment (suggestion):
`pandas` doesn't know what to do with these specs and from what I can tell, we never assert correctness on the result, so I've removed the tests and would consider this functionality not supported.

Code after:
    assert len(no_mean_finalizers) == len(no_mean_spec)
    assert len(with_mean_finalizers) == len(with_mean_spec)


@pytest.mark.parametrize("split_every", [1, 8])
@pytest.mark.parametrize("split_out", [2, 32])
def test_shuffle_aggregate(shuffle_method, split_out, split_every):
    pdf = pd.DataFrame(
        {
            "a": [1, 2, 3, 1, 1, 2, 4, 3, 7] * 100,
            "b": [4, 2, 7, 3, 3, 1, 1, 1, 2] * 100,
            "c": [0, 1, 2, 3, 4, 5, 6, 7, 8] * 100,
            "d": [3, 2, 1, 3, 2, 1, 2, 6, 4] * 100,
        },
        columns=["c", "b", "a", "d"],
    )
    ddf = dd.from_pandas(pdf, npartitions=100)

    spec = {"b": "mean", "c": ["min", "max"]}
    result = ddf.groupby(["a", "b"], sort=False).agg(
        spec,
        split_out=split_out,
        split_every=split_every,
        shuffle_method=shuffle_method,
    )
    expect = pdf.groupby(["a", "b"]).agg(spec)

    # Make sure "mean" dtype is consistent.
    # Pandas<1.3 will return int instead of float
    expect[("b", "mean")] = expect[("b", "mean")].astype(result[("b", "mean")].dtype)
    assert_eq(expect, result)


@pytest.mark.parametrize("sort", [True, False])
def test_shuffle_aggregate_sort(shuffle_method, sort):
    pdf = pd.DataFrame(
        {
            "a": [1, 2, 3, 1, 1, 2, 4, 3, 7] * 100,
            "b": [4, 2, 7, 3, 3, 1, 1, 1, 2] * 100,
            "c": [0, 1, 2, 3, 4, 5, 6, 7, 8] * 100,
            "d": [3, 2, 1, 3, 2, 1, 2, 6, 4] * 100,
        },
        columns=["c", "b", "a", "d"],
    )
    ddf = dd.from_pandas(pdf, npartitions=100)

    # Check single-column groupby
    spec = {"b": "mean", "c": ["min", "max"]}
    result = ddf.groupby("a", sort=sort).agg(
        spec, split_out=2, shuffle_backend=shuffle_method
    )