""" test all other .agg behavior """ import datetime as dt from functools import partial import numpy as np import pytest from pandas.errors import SpecificationError import pandas as pd from pandas import ( DataFrame, Index, MultiIndex, PeriodIndex, Series, date_range, period_range, ) import pandas._testing as tm from pandas.io.formats.printing import pprint_thing def test_agg_partial_failure_raises(): # GH#43741 df = DataFrame( { "data1": np.random.default_rng(2).standard_normal(5), "data2": np.random.default_rng(2).standard_normal(5), "key1": ["a", "a", "b", "b", "a"], "key2": ["one", "two", "one", "two", "one"], } ) grouped = df.groupby("key1") def peak_to_peak(arr): return arr.max() - arr.min() with pytest.raises(TypeError, match="unsupported operand type"): grouped.agg([peak_to_peak]) with pytest.raises(TypeError, match="unsupported operand type"): grouped.agg(peak_to_peak) def test_agg_datetimes_mixed(): data = [[1, "2012-01-01", 1.0], [2, "2012-01-02", 2.0], [3, None, 3.0]] df1 = DataFrame( { "key": [x[0] for x in data], "date": [x[1] for x in data], "value": [x[2] for x in data], } ) data = [ [ row[0], (dt.datetime.strptime(row[1], "%Y-%m-%d").date() if row[1] else None), row[2], ] for row in data ] df2 = DataFrame( { "key": [x[0] for x in data], "date": [x[1] for x in data], "value": [x[2] for x in data], } ) df1["weights"] = df1["value"] / df1["value"].sum() gb1 = df1.groupby("date").aggregate("sum") df2["weights"] = df1["value"] / df1["value"].sum() gb2 = df2.groupby("date").aggregate("sum") assert len(gb1) == len(gb2) def test_agg_period_index(): prng = period_range("2012-1-1", freq="M", periods=3) df = DataFrame(np.random.default_rng(2).standard_normal((3, 2)), index=prng) rs = df.groupby(level=0).sum() assert isinstance(rs.index, PeriodIndex) # GH 3579 index = period_range(start="1999-01", periods=5, freq="M") s1 = Series(np.random.default_rng(2).random(len(index)), index=index) s2 = Series(np.random.default_rng(2).random(len(index)), index=index) df = DataFrame.from_dict({"s1": s1, "s2": s2}) grouped = df.groupby(df.index.month) list(grouped) def test_agg_dict_parameter_cast_result_dtypes(): # GH 12821 df = DataFrame( { "class": ["A", "A", "B", "B", "C", "C", "D", "D"], "time": date_range("1/1/2011", periods=8, freq="H"), } ) df.loc[[0, 1, 2, 5], "time"] = None # test for `first` function exp = df.loc[[0, 3, 4, 6]].set_index("class") grouped = df.groupby("class") tm.assert_frame_equal(grouped.first(), exp) tm.assert_frame_equal(grouped.agg("first"), exp) tm.assert_frame_equal(grouped.agg({"time": "first"}), exp) tm.assert_series_equal(grouped.time.first(), exp["time"]) tm.assert_series_equal(grouped.time.agg("first"), exp["time"]) # test for `last` function exp = df.loc[[0, 3, 4, 7]].set_index("class") grouped = df.groupby("class") tm.assert_frame_equal(grouped.last(), exp) tm.assert_frame_equal(grouped.agg("last"), exp) tm.assert_frame_equal(grouped.agg({"time": "last"}), exp) tm.assert_series_equal(grouped.time.last(), exp["time"]) tm.assert_series_equal(grouped.time.agg("last"), exp["time"]) # count exp = Series([2, 2, 2, 2], index=Index(list("ABCD"), name="class"), name="time") tm.assert_series_equal(grouped.time.agg(len), exp) tm.assert_series_equal(grouped.time.size(), exp) exp = Series([0, 1, 1, 2], index=Index(list("ABCD"), name="class"), name="time") tm.assert_series_equal(grouped.time.count(), exp) def test_agg_cast_results_dtypes(): # similar to GH12821 # xref #11444 u = [dt.datetime(2015, x + 1, 1) for x in range(12)] v = list("aaabbbbbbccd") df = DataFrame({"X": v, "Y": u}) result = df.groupby("X")["Y"].agg(len) expected = df.groupby("X")["Y"].count() tm.assert_series_equal(result, expected) def test_aggregate_float64_no_int64(): # see gh-11199 df = DataFrame({"a": [1, 2, 3, 4, 5], "b": [1, 2, 2, 4, 5], "c": [1, 2, 3, 4, 5]}) expected = DataFrame({"a": [1, 2.5, 4, 5]}, index=[1, 2, 4, 5]) expected.index.name = "b" result = df.groupby("b")[["a"]].mean() tm.assert_frame_equal(result, expected) expected = DataFrame({"a": [1, 2.5, 4, 5], "c": [1, 2.5, 4, 5]}, index=[1, 2, 4, 5]) expected.index.name = "b" result = df.groupby("b")[["a", "c"]].mean() tm.assert_frame_equal(result, expected) def test_aggregate_api_consistency(): # GH 9052 # make sure that the aggregates via dict # are consistent df = DataFrame( { "A": ["foo", "bar", "foo", "bar", "foo", "bar", "foo", "foo"], "B": ["one", "one", "two", "two", "two", "two", "one", "two"], "C": np.random.default_rng(2).standard_normal(8) + 1.0, "D": np.arange(8), } ) grouped = df.groupby(["A", "B"]) c_mean = grouped["C"].mean() c_sum = grouped["C"].sum() d_mean = grouped["D"].mean() d_sum = grouped["D"].sum() result = grouped["D"].agg(["sum", "mean"]) expected = pd.concat([d_sum, d_mean], axis=1) expected.columns = ["sum", "mean"] tm.assert_frame_equal(result, expected, check_like=True) result = grouped.agg(["sum", "mean"]) expected = pd.concat([c_sum, c_mean, d_sum, d_mean], axis=1) expected.columns = MultiIndex.from_product([["C", "D"], ["sum", "mean"]]) tm.assert_frame_equal(result, expected, check_like=True) result = grouped[["D", "C"]].agg(["sum", "mean"]) expected = pd.concat([d_sum, d_mean, c_sum, c_mean], axis=1) expected.columns = MultiIndex.from_product([["D", "C"], ["sum", "mean"]]) tm.assert_frame_equal(result, expected, check_like=True) result = grouped.agg({"C": "mean", "D": "sum"}) expected = pd.concat([d_sum, c_mean], axis=1) tm.assert_frame_equal(result, expected, check_like=True) result = grouped.agg({"C": ["mean", "sum"], "D": ["mean", "sum"]}) expected = pd.concat([c_mean, c_sum, d_mean, d_sum], axis=1) expected.columns = MultiIndex.from_product([["C", "D"], ["mean", "sum"]]) msg = r"Column\(s\) \['r', 'r2'\] do not exist" with pytest.raises(KeyError, match=msg): grouped[["D", "C"]].agg({"r": "sum", "r2": "mean"}) def test_agg_dict_renaming_deprecation(): # 15931 df = DataFrame({"A": [1, 1, 1, 2, 2], "B": range(5), "C": range(5)}) msg = r"nested renamer is not supported" with pytest.raises(SpecificationError, match=msg): df.groupby("A").agg( {"B": {"foo": ["sum", "max"]}, "C": {"bar": ["count", "min"]}} ) msg = r"Column\(s\) \['ma'\] do not exist" with pytest.raises(KeyError, match=msg): df.groupby("A")[["B", "C"]].agg({"ma": "max"}) msg = r"nested renamer is not supported" with pytest.raises(SpecificationError, match=msg): df.groupby("A").B.agg({"foo": "count"}) def test_agg_compat(): # GH 12334 df = DataFrame( { "A": ["foo", "bar", "foo", "bar", "foo", "bar", "foo", "foo"], "B": ["one", "one", "two", "two", "two", "two", "one", "two"], "C": np.random.default_rng(2).standard_normal(8) + 1.0, "D": np.arange(8), } ) g = df.groupby(["A", "B"]) msg = r"nested renamer is not supported" with pytest.raises(SpecificationError, match=msg): g["D"].agg({"C": ["sum", "std"]}) with pytest.raises(SpecificationError, match=msg): g["D"].agg({"C": "sum", "D": "std"}) def test_agg_nested_dicts(): # API change for disallowing these types of nested dicts df = DataFrame( { "A": ["foo", "bar", "foo", "bar", "foo", "bar", "foo", "foo"], "B": ["one", "one", "two", "two", "two", "two", "one", "two"], "C": np.random.default_rng(2).standard_normal(8) + 1.0, "D": np.arange(8), } ) g = df.groupby(["A", "B"]) msg = r"nested renamer is not supported" with pytest.raises(SpecificationError, match=msg): g.aggregate({"r1": {"C": ["mean", "sum"]}, "r2": {"D": ["mean", "sum"]}}) with pytest.raises(SpecificationError, match=msg): g.agg({"C": {"ra": ["mean", "std"]}, "D": {"rb": ["mean", "std"]}}) # same name as the original column # GH9052 with pytest.raises(SpecificationError, match=msg): g["D"].agg({"result1": np.sum, "result2": np.mean}) with pytest.raises(SpecificationError, match=msg): g["D"].agg({"D": np.sum, "result2": np.mean}) def test_agg_item_by_item_raise_typeerror(): df = DataFrame(np.random.default_rng(2).integers(10, size=(20, 10))) def raiseException(df): pprint_thing("----------------------------------------") pprint_thing(df.to_string()) raise TypeError("test") with pytest.raises(TypeError, match="test"): df.groupby(0).agg(raiseException) def test_series_agg_multikey(): ts = tm.makeTimeSeries() grouped = ts.groupby([lambda x: x.year, lambda x: x.month]) result = grouped.agg("sum") expected = grouped.sum() tm.assert_series_equal(result, expected) def test_series_agg_multi_pure_python(): data = DataFrame( { "A": [ "foo", "foo", "foo", "foo", "bar", "bar", "bar", "bar", "foo", "foo", "foo", ], "B": [ "one", "one", "one", "two", "one", "one", "one", "two", "two", "two", "one", ], "C": [ "dull", "dull", "shiny", "dull", "dull", "shiny", "shiny", "dull", "shiny", "shiny", "shiny", ], "D": np.random.default_rng(2).standard_normal(11), "E": np.random.default_rng(2).standard_normal(11), "F": np.random.default_rng(2).standard_normal(11), } ) def bad(x): assert len(x.values.base) > 0 return "foo" result = data.groupby(["A", "B"]).agg(bad) expected = data.groupby(["A", "B"]).agg(lambda x: "foo") tm.assert_frame_equal(result, expected) def test_agg_consistency(): # agg with ([]) and () not consistent # GH 6715 def P1(a): return np.percentile(a.dropna(), q=1) df = DataFrame( { "col1": [1, 2, 3, 4], "col2": [10, 25, 26, 31], "date": [ dt.date(2013, 2, 10), dt.date(2013, 2, 10), dt.date(2013, 2, 11), dt.date(2013, 2, 11), ], } ) g = df.groupby("date") expected = g.agg([P1]) expected.columns = expected.columns.levels[0] result = g.agg(P1) tm.assert_frame_equal(result, expected) def test_agg_callables(): # GH 7929 df = DataFrame({"foo": [1, 2], "bar": [3, 4]}).astype(np.int64) class fn_class: def __call__(self, x): return sum(x) equiv_callables = [ sum, np.sum, lambda x: sum(x), lambda x: x.sum(), partial(sum), fn_class(), ] expected = df.groupby("foo").agg("sum") for ecall in equiv_callables: warn = FutureWarning if ecall is sum or ecall is np.sum else None msg = "using DataFrameGroupBy.sum" with tm.assert_produces_warning(warn, match=msg): result = df.groupby("foo").agg(ecall) tm.assert_frame_equal(result, expected) def test_agg_over_numpy_arrays(): # GH 3788 df = DataFrame( [ [1, np.array([10, 20, 30])], [1, np.array([40, 50, 60])], [2, np.array([20, 30, 40])], ], columns=["category", "arraydata"], ) gb = df.groupby("category") expected_data = [[np.array([50, 70, 90])], [np.array([20, 30, 40])]] expected_index = Index([1, 2], name="category") expected_column = ["arraydata"] expected = DataFrame(expected_data, index=expected_index, columns=expected_column) alt = gb.sum(numeric_only=False) tm.assert_frame_equal(alt, expected) result = gb.agg("sum", numeric_only=False) tm.assert_frame_equal(result, expected) # FIXME: the original version of this test called `gb.agg(sum)` # and that raises TypeError if `numeric_only=False` is passed @pytest.mark.parametrize("as_period", [True, False]) def test_agg_tzaware_non_datetime_result(as_period): # discussed in GH#29589, fixed in GH#29641, operating on tzaware values # with function that is not dtype-preserving dti = date_range("2012-01-01", periods=4, tz="UTC") if as_period: dti = dti.tz_localize(None).to_period("D") df = DataFrame({"a": [0, 0, 1, 1], "b": dti}) gb = df.groupby("a") # Case that _does_ preserve the dtype result = gb["b"].agg(lambda x: x.iloc[0]) expected = Series(dti[::2], name="b") expected.index.name = "a" tm.assert_series_equal(result, expected) # Cases that do _not_ preserve the dtype result = gb["b"].agg(lambda x: x.iloc[0].year) expected = Series([2012, 2012], name="b") expected.index.name = "a" tm.assert_series_equal(result, expected) result = gb["b"].agg(lambda x: x.iloc[-1] - x.iloc[0]) expected = Series([pd.Timedelta(days=1), pd.Timedelta(days=1)], name="b") expected.index.name = "a" if as_period: expected = Series([pd.offsets.Day(1), pd.offsets.Day(1)], name="b") expected.index.name = "a" tm.assert_series_equal(result, expected) def test_agg_timezone_round_trip(): # GH 15426 ts = pd.Timestamp("2016-01-01 12:00:00", tz="US/Pacific") df = DataFrame({"a": 1, "b": [ts + dt.timedelta(minutes=nn) for nn in range(10)]}) result1 = df.groupby("a")["b"].agg("min").iloc[0] result2 = df.groupby("a")["b"].agg(lambda x: np.min(x)).iloc[0] result3 = df.groupby("a")["b"].min().iloc[0] assert result1 == ts assert result2 == ts assert result3 == ts dates = [ pd.Timestamp(f"2016-01-0{i:d} 12:00:00", tz="US/Pacific") for i in range(1, 5) ] df = DataFrame({"A": ["a", "b"] * 2, "B": dates}) grouped = df.groupby("A") ts = df["B"].iloc[0] assert ts == grouped.nth(0)["B"].iloc[0] assert ts == grouped.head(1)["B"].iloc[0] assert ts == grouped.first()["B"].iloc[0] # GH#27110 applying iloc should return a DataFrame assert ts == grouped.apply(lambda x: x.iloc[0]).iloc[0, 1] ts = df["B"].iloc[2] assert ts == grouped.last()["B"].iloc[0] # GH#27110 applying iloc should return a DataFrame assert ts == grouped.apply(lambda x: x.iloc[-1]).iloc[0, 1] def test_sum_uint64_overflow(): # see gh-14758 # Convert to uint64 and don't overflow df = DataFrame([[1, 2], [3, 4], [5, 6]], dtype=object) df = df + 9223372036854775807 index = Index( [9223372036854775808, 9223372036854775810, 9223372036854775812], dtype=np.uint64 ) expected = DataFrame( {1: [9223372036854775809, 9223372036854775811, 9223372036854775813]}, index=index, dtype=object, ) expected.index.name = 0 result = df.groupby(0).sum(numeric_only=False) tm.assert_frame_equal(result, expected) # out column is non-numeric, so with numeric_only=True it is dropped result2 = df.groupby(0).sum(numeric_only=True) expected2 = expected[[]] tm.assert_frame_equal(result2, expected2) @pytest.mark.parametrize( "structure, expected", [ (tuple, DataFrame({"C": {(1, 1): (1, 1, 1), (3, 4): (3, 4, 4)}})), (list, DataFrame({"C": {(1, 1): [1, 1, 1], (3, 4): [3, 4, 4]}})), ( lambda x: tuple(x), DataFrame({"C": {(1, 1): (1, 1, 1), (3, 4): (3, 4, 4)}}), ), ( lambda x: list(x), DataFrame({"C": {(1, 1): [1, 1, 1], (3, 4): [3, 4, 4]}}), ), ], ) def test_agg_structs_dataframe(structure, expected): df = DataFrame( {"A": [1, 1, 1, 3, 3, 3], "B": [1, 1, 1, 4, 4, 4], "C": [1, 1, 1, 3, 4, 4]} ) result = df.groupby(["A", "B"]).aggregate(structure) expected.index.names = ["A", "B"] tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "structure, expected", [ (tuple, Series([(1, 1, 1), (3, 4, 4)], index=[1, 3], name="C")), (list, Series([[1, 1, 1], [3, 4, 4]], index=[1, 3], name="C")), (lambda x: tuple(x), Series([(1, 1, 1), (3, 4, 4)], index=[1, 3], name="C")), (lambda x: list(x), Series([[1, 1, 1], [3, 4, 4]], index=[1, 3], name="C")), ], ) def test_agg_structs_series(structure, expected): # Issue #18079 df = DataFrame( {"A": [1, 1, 1, 3, 3, 3], "B": [1, 1, 1, 4, 4, 4], "C": [1, 1, 1, 3, 4, 4]} ) result = df.groupby("A")["C"].aggregate(structure) expected.index.name = "A" tm.assert_series_equal(result, expected) def test_agg_category_nansum(observed): categories = ["a", "b", "c"] df = DataFrame( {"A": pd.Categorical(["a", "a", "b"], categories=categories), "B": [1, 2, 3]} ) msg = "using SeriesGroupBy.sum" with tm.assert_produces_warning(FutureWarning, match=msg): result = df.groupby("A", observed=observed).B.agg(np.nansum) expected = Series( [3, 3, 0], index=pd.CategoricalIndex(["a", "b", "c"], categories=categories, name="A"), name="B", ) if observed: expected = expected[expected != 0] tm.assert_series_equal(result, expected) def test_agg_list_like_func(): # GH 18473 df = DataFrame({"A": [str(x) for x in range(3)], "B": [str(x) for x in range(3)]}) grouped = df.groupby("A", as_index=False, sort=False) result = grouped.agg({"B": lambda x: list(x)}) expected = DataFrame( {"A": [str(x) for x in range(3)], "B": [[str(x)] for x in range(3)]} ) tm.assert_frame_equal(result, expected) def test_agg_lambda_with_timezone(): # GH 23683 df = DataFrame( { "tag": [1, 1], "date": [ pd.Timestamp("2018-01-01", tz="UTC"), pd.Timestamp("2018-01-02", tz="UTC"), ], } ) result = df.groupby("tag").agg({"date": lambda e: e.head(1)}) expected = DataFrame( [pd.Timestamp("2018-01-01", tz="UTC")], index=Index([1], name="tag"), columns=["date"], ) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "err_cls", [ NotImplementedError, RuntimeError, KeyError, IndexError, OSError, ValueError, ArithmeticError, AttributeError, ], ) def test_groupby_agg_err_catching(err_cls): # make sure we suppress anything other than TypeError or AssertionError # in _python_agg_general # Use a non-standard EA to make sure we don't go down ndarray paths from pandas.tests.extension.decimal.array import ( DecimalArray, make_data, to_decimal, ) data = make_data()[:5] df = DataFrame( {"id1": [0, 0, 0, 1, 1], "id2": [0, 1, 0, 1, 1], "decimals": DecimalArray(data)} ) expected = Series(to_decimal([data[0], data[3]])) def weird_func(x): # weird function that raise something other than TypeError or IndexError # in _python_agg_general if len(x) == 0: raise err_cls return x.iloc[0] result = df["decimals"].groupby(df["id1"]).agg(weird_func) tm.assert_series_equal(result, expected, check_names=False)