123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482 |
- """ test with the .transform """
- from io import StringIO
- import numpy as np
- import pytest
- from pandas.core.dtypes.common import (
- ensure_platform_int,
- is_timedelta64_dtype,
- )
- import pandas as pd
- from pandas import (
- Categorical,
- DataFrame,
- MultiIndex,
- Series,
- Timestamp,
- concat,
- date_range,
- )
- import pandas._testing as tm
- from pandas.tests.groupby import get_groupby_method_args
- def assert_fp_equal(a, b):
- assert (np.abs(a - b) < 1e-12).all()
- def test_transform():
- data = Series(np.arange(9) // 3, index=np.arange(9))
- index = np.arange(9)
- np.random.shuffle(index)
- data = data.reindex(index)
- grouped = data.groupby(lambda x: x // 3)
- transformed = grouped.transform(lambda x: x * x.sum())
- assert transformed[7] == 12
- # GH 8046
- # make sure that we preserve the input order
- df = DataFrame(
- np.arange(6, dtype="int64").reshape(3, 2), columns=["a", "b"], index=[0, 2, 1]
- )
- key = [0, 0, 1]
- expected = (
- df.sort_index()
- .groupby(key)
- .transform(lambda x: x - x.mean())
- .groupby(key)
- .mean()
- )
- result = df.groupby(key).transform(lambda x: x - x.mean()).groupby(key).mean()
- tm.assert_frame_equal(result, expected)
- def demean(arr):
- return arr - arr.mean(axis=0)
- people = DataFrame(
- np.random.randn(5, 5),
- columns=["a", "b", "c", "d", "e"],
- index=["Joe", "Steve", "Wes", "Jim", "Travis"],
- )
- key = ["one", "two", "one", "two", "one"]
- result = people.groupby(key).transform(demean).groupby(key).mean()
- expected = people.groupby(key, group_keys=False).apply(demean).groupby(key).mean()
- tm.assert_frame_equal(result, expected)
- # GH 8430
- df = tm.makeTimeDataFrame()
- g = df.groupby(pd.Grouper(freq="M"))
- g.transform(lambda x: x - 1)
- # GH 9700
- df = DataFrame({"a": range(5, 10), "b": range(5)})
- result = df.groupby("a").transform(max)
- expected = DataFrame({"b": range(5)})
- tm.assert_frame_equal(result, expected)
- def test_transform_fast():
- df = DataFrame({"id": np.arange(100000) / 3, "val": np.random.randn(100000)})
- grp = df.groupby("id")["val"]
- values = np.repeat(grp.mean().values, ensure_platform_int(grp.count().values))
- expected = Series(values, index=df.index, name="val")
- result = grp.transform(np.mean)
- tm.assert_series_equal(result, expected)
- result = grp.transform("mean")
- tm.assert_series_equal(result, expected)
- # GH 12737
- df = DataFrame(
- {
- "grouping": [0, 1, 1, 3],
- "f": [1.1, 2.1, 3.1, 4.5],
- "d": date_range("2014-1-1", "2014-1-4"),
- "i": [1, 2, 3, 4],
- },
- columns=["grouping", "f", "i", "d"],
- )
- result = df.groupby("grouping").transform("first")
- dates = [
- Timestamp("2014-1-1"),
- Timestamp("2014-1-2"),
- Timestamp("2014-1-2"),
- Timestamp("2014-1-4"),
- ]
- expected = DataFrame(
- {"f": [1.1, 2.1, 2.1, 4.5], "d": dates, "i": [1, 2, 2, 4]},
- columns=["f", "i", "d"],
- )
- tm.assert_frame_equal(result, expected)
- # selection
- result = df.groupby("grouping")[["f", "i"]].transform("first")
- expected = expected[["f", "i"]]
- tm.assert_frame_equal(result, expected)
- # dup columns
- df = DataFrame([[1, 2, 3], [4, 5, 6]], columns=["g", "a", "a"])
- result = df.groupby("g").transform("first")
- expected = df.drop("g", axis=1)
- tm.assert_frame_equal(result, expected)
- def test_transform_broadcast(tsframe, ts):
- grouped = ts.groupby(lambda x: x.month)
- result = grouped.transform(np.mean)
- tm.assert_index_equal(result.index, ts.index)
- for _, gp in grouped:
- assert_fp_equal(result.reindex(gp.index), gp.mean())
- grouped = tsframe.groupby(lambda x: x.month)
- result = grouped.transform(np.mean)
- tm.assert_index_equal(result.index, tsframe.index)
- for _, gp in grouped:
- agged = gp.mean(axis=0)
- res = result.reindex(gp.index)
- for col in tsframe:
- assert_fp_equal(res[col], agged[col])
- # group columns
- grouped = tsframe.groupby({"A": 0, "B": 0, "C": 1, "D": 1}, axis=1)
- result = grouped.transform(np.mean)
- tm.assert_index_equal(result.index, tsframe.index)
- tm.assert_index_equal(result.columns, tsframe.columns)
- for _, gp in grouped:
- agged = gp.mean(1)
- res = result.reindex(columns=gp.columns)
- for idx in gp.index:
- assert_fp_equal(res.xs(idx), agged[idx])
- def test_transform_axis_1(request, transformation_func):
- # GH 36308
- df = DataFrame({"a": [1, 2], "b": [3, 4], "c": [5, 6]}, index=["x", "y"])
- args = get_groupby_method_args(transformation_func, df)
- result = df.groupby([0, 0, 1], axis=1).transform(transformation_func, *args)
- expected = df.T.groupby([0, 0, 1]).transform(transformation_func, *args).T
- if transformation_func in ["diff", "shift"]:
- # Result contains nans, so transpose coerces to float
- expected["b"] = expected["b"].astype("int64")
- # cumcount returns Series; the rest are DataFrame
- tm.assert_equal(result, expected)
- def test_transform_axis_1_reducer(request, reduction_func):
- # GH#45715
- if reduction_func in (
- "corrwith",
- "ngroup",
- "nth",
- ):
- marker = pytest.mark.xfail(reason="transform incorrectly fails - GH#45986")
- request.node.add_marker(marker)
- df = DataFrame({"a": [1, 2], "b": [3, 4], "c": [5, 6]}, index=["x", "y"])
- result = df.groupby([0, 0, 1], axis=1).transform(reduction_func)
- expected = df.T.groupby([0, 0, 1]).transform(reduction_func).T
- tm.assert_equal(result, expected)
- def test_transform_axis_ts(tsframe):
- # make sure that we are setting the axes
- # correctly when on axis=0 or 1
- # in the presence of a non-monotonic indexer
- # GH12713
- base = tsframe.iloc[0:5]
- r = len(base.index)
- c = len(base.columns)
- tso = DataFrame(
- np.random.randn(r, c), index=base.index, columns=base.columns, dtype="float64"
- )
- # monotonic
- ts = tso
- grouped = ts.groupby(lambda x: x.weekday(), group_keys=False)
- result = ts - grouped.transform("mean")
- expected = grouped.apply(lambda x: x - x.mean(axis=0))
- tm.assert_frame_equal(result, expected)
- ts = ts.T
- grouped = ts.groupby(lambda x: x.weekday(), axis=1, group_keys=False)
- result = ts - grouped.transform("mean")
- expected = grouped.apply(lambda x: (x.T - x.mean(1)).T)
- tm.assert_frame_equal(result, expected)
- # non-monotonic
- ts = tso.iloc[[1, 0] + list(range(2, len(base)))]
- grouped = ts.groupby(lambda x: x.weekday(), group_keys=False)
- result = ts - grouped.transform("mean")
- expected = grouped.apply(lambda x: x - x.mean(axis=0))
- tm.assert_frame_equal(result, expected)
- ts = ts.T
- grouped = ts.groupby(lambda x: x.weekday(), axis=1, group_keys=False)
- result = ts - grouped.transform("mean")
- expected = grouped.apply(lambda x: (x.T - x.mean(1)).T)
- tm.assert_frame_equal(result, expected)
- def test_transform_dtype():
- # GH 9807
- # Check transform dtype output is preserved
- df = DataFrame([[1, 3], [2, 3]])
- result = df.groupby(1).transform("mean")
- expected = DataFrame([[1.5], [1.5]])
- tm.assert_frame_equal(result, expected)
- def test_transform_bug():
- # GH 5712
- # transforming on a datetime column
- df = DataFrame({"A": Timestamp("20130101"), "B": np.arange(5)})
- result = df.groupby("A")["B"].transform(lambda x: x.rank(ascending=False))
- expected = Series(np.arange(5, 0, step=-1), name="B", dtype="float64")
- tm.assert_series_equal(result, expected)
- def test_transform_numeric_to_boolean():
- # GH 16875
- # inconsistency in transforming boolean values
- expected = Series([True, True], name="A")
- df = DataFrame({"A": [1.1, 2.2], "B": [1, 2]})
- result = df.groupby("B").A.transform(lambda x: True)
- tm.assert_series_equal(result, expected)
- df = DataFrame({"A": [1, 2], "B": [1, 2]})
- result = df.groupby("B").A.transform(lambda x: True)
- tm.assert_series_equal(result, expected)
- def test_transform_datetime_to_timedelta():
- # GH 15429
- # transforming a datetime to timedelta
- df = DataFrame({"A": Timestamp("20130101"), "B": np.arange(5)})
- expected = Series([Timestamp("20130101") - Timestamp("20130101")] * 5, name="A")
- # this does date math without changing result type in transform
- base_time = df["A"][0]
- result = (
- df.groupby("A")["A"].transform(lambda x: x.max() - x.min() + base_time)
- - base_time
- )
- tm.assert_series_equal(result, expected)
- # this does date math and causes the transform to return timedelta
- result = df.groupby("A")["A"].transform(lambda x: x.max() - x.min())
- tm.assert_series_equal(result, expected)
- def test_transform_datetime_to_numeric():
- # GH 10972
- # convert dt to float
- df = DataFrame({"a": 1, "b": date_range("2015-01-01", periods=2, freq="D")})
- result = df.groupby("a").b.transform(
- lambda x: x.dt.dayofweek - x.dt.dayofweek.mean()
- )
- expected = Series([-0.5, 0.5], name="b")
- tm.assert_series_equal(result, expected)
- # convert dt to int
- df = DataFrame({"a": 1, "b": date_range("2015-01-01", periods=2, freq="D")})
- result = df.groupby("a").b.transform(
- lambda x: x.dt.dayofweek - x.dt.dayofweek.min()
- )
- expected = Series([0, 1], dtype=np.int32, name="b")
- tm.assert_series_equal(result, expected)
- def test_transform_casting():
- # 13046
- data = """
- idx A ID3 DATETIME
- 0 B-028 b76cd912ff "2014-10-08 13:43:27"
- 1 B-054 4a57ed0b02 "2014-10-08 14:26:19"
- 2 B-076 1a682034f8 "2014-10-08 14:29:01"
- 3 B-023 b76cd912ff "2014-10-08 18:39:34"
- 4 B-023 f88g8d7sds "2014-10-08 18:40:18"
- 5 B-033 b76cd912ff "2014-10-08 18:44:30"
- 6 B-032 b76cd912ff "2014-10-08 18:46:00"
- 7 B-037 b76cd912ff "2014-10-08 18:52:15"
- 8 B-046 db959faf02 "2014-10-08 18:59:59"
- 9 B-053 b76cd912ff "2014-10-08 19:17:48"
- 10 B-065 b76cd912ff "2014-10-08 19:21:38"
- """
- df = pd.read_csv(
- StringIO(data), sep=r"\s+", index_col=[0], parse_dates=["DATETIME"]
- )
- result = df.groupby("ID3")["DATETIME"].transform(lambda x: x.diff())
- assert is_timedelta64_dtype(result.dtype)
- result = df[["ID3", "DATETIME"]].groupby("ID3").transform(lambda x: x.diff())
- assert is_timedelta64_dtype(result.DATETIME.dtype)
- def test_transform_multiple(ts):
- grouped = ts.groupby([lambda x: x.year, lambda x: x.month])
- grouped.transform(lambda x: x * 2)
- grouped.transform(np.mean)
- def test_dispatch_transform(tsframe):
- df = tsframe[::5].reindex(tsframe.index)
- grouped = df.groupby(lambda x: x.month)
- filled = grouped.fillna(method="pad")
- fillit = lambda x: x.fillna(method="pad")
- expected = df.groupby(lambda x: x.month).transform(fillit)
- tm.assert_frame_equal(filled, expected)
- def test_transform_transformation_func(transformation_func):
- # GH 30918
- df = DataFrame(
- {
- "A": ["foo", "foo", "foo", "foo", "bar", "bar", "baz"],
- "B": [1, 2, np.nan, 3, 3, np.nan, 4],
- },
- index=date_range("2020-01-01", "2020-01-07"),
- )
- if transformation_func == "cumcount":
- test_op = lambda x: x.transform("cumcount")
- mock_op = lambda x: Series(range(len(x)), x.index)
- elif transformation_func == "fillna":
- test_op = lambda x: x.transform("fillna", value=0)
- mock_op = lambda x: x.fillna(value=0)
- elif transformation_func == "ngroup":
- test_op = lambda x: x.transform("ngroup")
- counter = -1
- def mock_op(x):
- nonlocal counter
- counter += 1
- return Series(counter, index=x.index)
- else:
- test_op = lambda x: x.transform(transformation_func)
- mock_op = lambda x: getattr(x, transformation_func)()
- result = test_op(df.groupby("A"))
- # pass the group in same order as iterating `for ... in df.groupby(...)`
- # but reorder to match df's index since this is a transform
- groups = [df[["B"]].iloc[4:6], df[["B"]].iloc[6:], df[["B"]].iloc[:4]]
- expected = concat([mock_op(g) for g in groups]).sort_index()
- # sort_index does not preserve the freq
- expected = expected.set_axis(df.index)
- if transformation_func in ("cumcount", "ngroup"):
- tm.assert_series_equal(result, expected)
- else:
- tm.assert_frame_equal(result, expected)
- def test_transform_select_columns(df):
- f = lambda x: x.mean()
- result = df.groupby("A")[["C", "D"]].transform(f)
- selection = df[["C", "D"]]
- expected = selection.groupby(df["A"]).transform(f)
- tm.assert_frame_equal(result, expected)
- def test_transform_nuisance_raises(df):
- # case that goes through _transform_item_by_item
- df.columns = ["A", "B", "B", "D"]
- # this also tests orderings in transform between
- # series/frame to make sure it's consistent
- grouped = df.groupby("A")
- gbc = grouped["B"]
- with pytest.raises(TypeError, match="Could not convert"):
- gbc.transform(lambda x: np.mean(x))
- with pytest.raises(TypeError, match="Could not convert"):
- df.groupby("A").transform(lambda x: np.mean(x))
- def test_transform_function_aliases(df):
- result = df.groupby("A").transform("mean", numeric_only=True)
- expected = df.groupby("A")[["C", "D"]].transform(np.mean)
- tm.assert_frame_equal(result, expected)
- result = df.groupby("A")["C"].transform("mean")
- expected = df.groupby("A")["C"].transform(np.mean)
- tm.assert_series_equal(result, expected)
- def test_series_fast_transform_date():
- # GH 13191
- df = DataFrame(
- {"grouping": [np.nan, 1, 1, 3], "d": date_range("2014-1-1", "2014-1-4")}
- )
- result = df.groupby("grouping")["d"].transform("first")
- dates = [
- pd.NaT,
- Timestamp("2014-1-2"),
- Timestamp("2014-1-2"),
- Timestamp("2014-1-4"),
- ]
- expected = Series(dates, name="d")
- tm.assert_series_equal(result, expected)
- def test_transform_length():
- # GH 9697
- df = DataFrame({"col1": [1, 1, 2, 2], "col2": [1, 2, 3, np.nan]})
- expected = Series([3.0] * 4)
- def nsum(x):
- return np.nansum(x)
- results = [
- df.groupby("col1").transform(sum)["col2"],
- df.groupby("col1")["col2"].transform(sum),
- df.groupby("col1").transform(nsum)["col2"],
- df.groupby("col1")["col2"].transform(nsum),
- ]
- for result in results:
- tm.assert_series_equal(result, expected, check_names=False)
- def test_transform_coercion():
- # 14457
- # when we are transforming be sure to not coerce
- # via assignment
- df = DataFrame({"A": ["a", "a", "b", "b"], "B": [0, 1, 3, 4]})
- g = df.groupby("A")
- expected = g.transform(np.mean)
- result = g.transform(lambda x: np.mean(x, axis=0))
- tm.assert_frame_equal(result, expected)
- def test_groupby_transform_with_int():
- # GH 3740, make sure that we might upcast on item-by-item transform
- # floats
- df = DataFrame(
- {
- "A": [1, 1, 1, 2, 2, 2],
- "B": Series(1, dtype="float64"),
- "C": Series([1, 2, 3, 1, 2, 3], dtype="float64"),
- "D": "foo",
- }
- )
- with np.errstate(all="ignore"):
- result = df.groupby("A")[["B", "C"]].transform(
- lambda x: (x - x.mean()) / x.std()
- )
- expected = DataFrame(
- {"B": np.nan, "C": Series([-1, 0, 1, -1, 0, 1], dtype="float64")}
- )
- tm.assert_frame_equal(result, expected)
- # int case
- df = DataFrame(
- {
- "A": [1, 1, 1, 2, 2, 2],
- "B": 1,
- "C": [1, 2, 3, 1, 2, 3],
- "D": "foo",
- }
- )
- with np.errstate(all="ignore"):
- with pytest.raises(TypeError, match="Could not convert"):
- df.groupby("A").transform(lambda x: (x - x.mean()) / x.std())
- result = df.groupby("A")[["B", "C"]].transform(
- lambda x: (x - x.mean()) / x.std()
- )
- expected = DataFrame({"B": np.nan, "C": [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0]})
- tm.assert_frame_equal(result, expected)
- # int that needs float conversion
- s = Series([2, 3, 4, 10, 5, -1])
- df = DataFrame({"A": [1, 1, 1, 2, 2, 2], "B": 1, "C": s, "D": "foo"})
- with np.errstate(all="ignore"):
- with pytest.raises(TypeError, match="Could not convert"):
- df.groupby("A").transform(lambda x: (x - x.mean()) / x.std())
- result = df.groupby("A")[["B", "C"]].transform(
- lambda x: (x - x.mean()) / x.std()
- )
- s1 = s.iloc[0:3]
- s1 = (s1 - s1.mean()) / s1.std()
- s2 = s.iloc[3:6]
- s2 = (s2 - s2.mean()) / s2.std()
- expected = DataFrame({"B": np.nan, "C": concat([s1, s2])})
- tm.assert_frame_equal(result, expected)
- # int doesn't get downcasted
- result = df.groupby("A")[["B", "C"]].transform(lambda x: x * 2 / 2)
- expected = DataFrame({"B": 1.0, "C": [2.0, 3.0, 4.0, 10.0, 5.0, -1.0]})
- tm.assert_frame_equal(result, expected)
- def test_groupby_transform_with_nan_group():
- # GH 9941
- df = DataFrame({"a": range(10), "b": [1, 1, 2, 3, np.nan, 4, 4, 5, 5, 5]})
- result = df.groupby(df.b)["a"].transform(max)
- expected = Series([1.0, 1.0, 2.0, 3.0, np.nan, 6.0, 6.0, 9.0, 9.0, 9.0], name="a")
- tm.assert_series_equal(result, expected)
- def test_transform_mixed_type():
- index = MultiIndex.from_arrays([[0, 0, 0, 1, 1, 1], [1, 2, 3, 1, 2, 3]])
- df = DataFrame(
- {
- "d": [1.0, 1.0, 1.0, 2.0, 2.0, 2.0],
- "c": np.tile(["a", "b", "c"], 2),
- "v": np.arange(1.0, 7.0),
- },
- index=index,
- )
- def f(group):
- group["g"] = group["d"] * 2
- return group[:1]
- grouped = df.groupby("c")
- result = grouped.apply(f)
- assert result["d"].dtype == np.float64
- # this is by definition a mutating operation!
- with pd.option_context("mode.chained_assignment", None):
- for key, group in grouped:
- res = f(group)
- tm.assert_frame_equal(res, result.loc[key])
- @pytest.mark.parametrize(
- "op, args, targop",
- [
- ("cumprod", (), lambda x: x.cumprod()),
- ("cumsum", (), lambda x: x.cumsum()),
- ("shift", (-1,), lambda x: x.shift(-1)),
- ("shift", (1,), lambda x: x.shift()),
- ],
- )
- def test_cython_transform_series(op, args, targop):
- # GH 4095
- s = Series(np.random.randn(1000))
- s_missing = s.copy()
- s_missing.iloc[2:10] = np.nan
- labels = np.random.randint(0, 50, size=1000).astype(float)
- # series
- for data in [s, s_missing]:
- # print(data.head())
- expected = data.groupby(labels).transform(targop)
- tm.assert_series_equal(expected, data.groupby(labels).transform(op, *args))
- tm.assert_series_equal(expected, getattr(data.groupby(labels), op)(*args))
- @pytest.mark.parametrize("op", ["cumprod", "cumsum"])
- @pytest.mark.parametrize("skipna", [False, True])
- @pytest.mark.parametrize(
- "input, exp",
- [
- # When everything is NaN
- ({"key": ["b"] * 10, "value": np.nan}, Series([np.nan] * 10, name="value")),
- # When there is a single NaN
- (
- {"key": ["b"] * 10 + ["a"] * 2, "value": [3] * 3 + [np.nan] + [3] * 8},
- {
- ("cumprod", False): [3.0, 9.0, 27.0] + [np.nan] * 7 + [3.0, 9.0],
- ("cumprod", True): [
- 3.0,
- 9.0,
- 27.0,
- np.nan,
- 81.0,
- 243.0,
- 729.0,
- 2187.0,
- 6561.0,
- 19683.0,
- 3.0,
- 9.0,
- ],
- ("cumsum", False): [3.0, 6.0, 9.0] + [np.nan] * 7 + [3.0, 6.0],
- ("cumsum", True): [
- 3.0,
- 6.0,
- 9.0,
- np.nan,
- 12.0,
- 15.0,
- 18.0,
- 21.0,
- 24.0,
- 27.0,
- 3.0,
- 6.0,
- ],
- },
- ),
- ],
- )
- def test_groupby_cum_skipna(op, skipna, input, exp):
- df = DataFrame(input)
- result = df.groupby("key")["value"].transform(op, skipna=skipna)
- if isinstance(exp, dict):
- expected = exp[(op, skipna)]
- else:
- expected = exp
- expected = Series(expected, name="value")
- tm.assert_series_equal(expected, result)
- @pytest.mark.slow
- @pytest.mark.parametrize(
- "op, args, targop",
- [
- ("cumprod", (), lambda x: x.cumprod()),
- ("cumsum", (), lambda x: x.cumsum()),
- ("shift", (-1,), lambda x: x.shift(-1)),
- ("shift", (1,), lambda x: x.shift()),
- ],
- )
- def test_cython_transform_frame(op, args, targop):
- s = Series(np.random.randn(1000))
- s_missing = s.copy()
- s_missing.iloc[2:10] = np.nan
- labels = np.random.randint(0, 50, size=1000).astype(float)
- strings = list("qwertyuiopasdfghjklz")
- strings_missing = strings[:]
- strings_missing[5] = np.nan
- df = DataFrame(
- {
- "float": s,
- "float_missing": s_missing,
- "int": [1, 1, 1, 1, 2] * 200,
- "datetime": date_range("1990-1-1", periods=1000),
- "timedelta": pd.timedelta_range(1, freq="s", periods=1000),
- "string": strings * 50,
- "string_missing": strings_missing * 50,
- },
- columns=[
- "float",
- "float_missing",
- "int",
- "datetime",
- "timedelta",
- "string",
- "string_missing",
- ],
- )
- df["cat"] = df["string"].astype("category")
- df2 = df.copy()
- df2.index = MultiIndex.from_product([range(100), range(10)])
- # DataFrame - Single and MultiIndex,
- # group by values, index level, columns
- for df in [df, df2]:
- for gb_target in [
- {"by": labels},
- {"level": 0},
- {"by": "string"},
- ]: # {"by": 'string_missing'}]:
- # {"by": ['int','string']}]:
- # TODO: remove or enable commented-out code
- gb = df.groupby(group_keys=False, **gb_target)
- if op != "shift" and "int" not in gb_target:
- # numeric apply fastpath promotes dtype so have
- # to apply separately and concat
- i = gb[["int"]].apply(targop)
- f = gb[["float", "float_missing"]].apply(targop)
- expected = concat([f, i], axis=1)
- else:
- expected = gb.apply(targop)
- expected = expected.sort_index(axis=1)
- result = gb[expected.columns].transform(op, *args).sort_index(axis=1)
- tm.assert_frame_equal(result, expected)
- result = getattr(gb[expected.columns], op)(*args).sort_index(axis=1)
- tm.assert_frame_equal(result, expected)
- # individual columns
- for c in df:
- if (
- c not in ["float", "int", "float_missing"]
- and op != "shift"
- and not (c == "timedelta" and op == "cumsum")
- ):
- msg = "|".join(
- [
- "does not support .* operations",
- ".* is not supported for object dtype",
- "is not implemented for this dtype",
- ]
- )
- with pytest.raises(TypeError, match=msg):
- gb[c].transform(op)
- with pytest.raises(TypeError, match=msg):
- getattr(gb[c], op)()
- else:
- expected = gb[c].apply(targop)
- expected.name = c
- tm.assert_series_equal(expected, gb[c].transform(op, *args))
- tm.assert_series_equal(expected, getattr(gb[c], op)(*args))
- def test_transform_with_non_scalar_group():
- # GH 10165
- cols = MultiIndex.from_tuples(
- [
- ("syn", "A"),
- ("mis", "A"),
- ("non", "A"),
- ("syn", "C"),
- ("mis", "C"),
- ("non", "C"),
- ("syn", "T"),
- ("mis", "T"),
- ("non", "T"),
- ("syn", "G"),
- ("mis", "G"),
- ("non", "G"),
- ]
- )
- df = DataFrame(
- np.random.randint(1, 10, (4, 12)), columns=cols, index=["A", "C", "G", "T"]
- )
- msg = "transform must return a scalar value for each group.*"
- with pytest.raises(ValueError, match=msg):
- df.groupby(axis=1, level=1).transform(lambda z: z.div(z.sum(axis=1), axis=0))
- @pytest.mark.parametrize(
- "cols,expected",
- [
- ("a", Series([1, 1, 1], name="a")),
- (
- ["a", "c"],
- DataFrame({"a": [1, 1, 1], "c": [1, 1, 1]}),
- ),
- ],
- )
- @pytest.mark.parametrize("agg_func", ["count", "rank", "size"])
- def test_transform_numeric_ret(cols, expected, agg_func):
- # GH#19200 and GH#27469
- df = DataFrame(
- {"a": date_range("2018-01-01", periods=3), "b": range(3), "c": range(7, 10)}
- )
- result = df.groupby("b")[cols].transform(agg_func)
- if agg_func == "rank":
- expected = expected.astype("float")
- elif agg_func == "size" and cols == ["a", "c"]:
- # transform("size") returns a Series
- expected = expected["a"].rename(None)
- tm.assert_equal(result, expected)
- def test_transform_ffill():
- # GH 24211
- data = [["a", 0.0], ["a", float("nan")], ["b", 1.0], ["b", float("nan")]]
- df = DataFrame(data, columns=["key", "values"])
- result = df.groupby("key").transform("ffill")
- expected = DataFrame({"values": [0.0, 0.0, 1.0, 1.0]})
- tm.assert_frame_equal(result, expected)
- result = df.groupby("key")["values"].transform("ffill")
- expected = Series([0.0, 0.0, 1.0, 1.0], name="values")
- tm.assert_series_equal(result, expected)
- @pytest.mark.parametrize("mix_groupings", [True, False])
- @pytest.mark.parametrize("as_series", [True, False])
- @pytest.mark.parametrize("val1,val2", [("foo", "bar"), (1, 2), (1.0, 2.0)])
- @pytest.mark.parametrize(
- "fill_method,limit,exp_vals",
- [
- (
- "ffill",
- None,
- [np.nan, np.nan, "val1", "val1", "val1", "val2", "val2", "val2"],
- ),
- ("ffill", 1, [np.nan, np.nan, "val1", "val1", np.nan, "val2", "val2", np.nan]),
- (
- "bfill",
- None,
- ["val1", "val1", "val1", "val2", "val2", "val2", np.nan, np.nan],
- ),
- ("bfill", 1, [np.nan, "val1", "val1", np.nan, "val2", "val2", np.nan, np.nan]),
- ],
- )
- def test_group_fill_methods(
- mix_groupings, as_series, val1, val2, fill_method, limit, exp_vals
- ):
- vals = [np.nan, np.nan, val1, np.nan, np.nan, val2, np.nan, np.nan]
- _exp_vals = list(exp_vals)
- # Overwrite placeholder values
- for index, exp_val in enumerate(_exp_vals):
- if exp_val == "val1":
- _exp_vals[index] = val1
- elif exp_val == "val2":
- _exp_vals[index] = val2
- # Need to modify values and expectations depending on the
- # Series / DataFrame that we ultimately want to generate
- if mix_groupings: # ['a', 'b', 'a, 'b', ...]
- keys = ["a", "b"] * len(vals)
- def interweave(list_obj):
- temp = []
- for x in list_obj:
- temp.extend([x, x])
- return temp
- _exp_vals = interweave(_exp_vals)
- vals = interweave(vals)
- else: # ['a', 'a', 'a', ... 'b', 'b', 'b']
- keys = ["a"] * len(vals) + ["b"] * len(vals)
- _exp_vals = _exp_vals * 2
- vals = vals * 2
- df = DataFrame({"key": keys, "val": vals})
- if as_series:
- result = getattr(df.groupby("key")["val"], fill_method)(limit=limit)
- exp = Series(_exp_vals, name="val")
- tm.assert_series_equal(result, exp)
- else:
- result = getattr(df.groupby("key"), fill_method)(limit=limit)
- exp = DataFrame({"val": _exp_vals})
- tm.assert_frame_equal(result, exp)
- @pytest.mark.parametrize("fill_method", ["ffill", "bfill"])
- def test_pad_stable_sorting(fill_method):
- # GH 21207
- x = [0] * 20
- y = [np.nan] * 10 + [1] * 10
- if fill_method == "bfill":
- y = y[::-1]
- df = DataFrame({"x": x, "y": y})
- expected = df.drop("x", axis=1)
- result = getattr(df.groupby("x"), fill_method)()
- tm.assert_frame_equal(result, expected)
- @pytest.mark.parametrize(
- "freq",
- [
- None,
- pytest.param(
- "D",
- marks=pytest.mark.xfail(
- reason="GH#23918 before method uses freq in vectorized approach"
- ),
- ),
- ],
- )
- @pytest.mark.parametrize("periods", [1, -1])
- @pytest.mark.parametrize("fill_method", ["ffill", "bfill", None])
- @pytest.mark.parametrize("limit", [None, 1])
- def test_pct_change(frame_or_series, freq, periods, fill_method, limit):
- # GH 21200, 21621, 30463
- vals = [3, np.nan, np.nan, np.nan, 1, 2, 4, 10, np.nan, 4]
- keys = ["a", "b"]
- key_v = np.repeat(keys, len(vals))
- df = DataFrame({"key": key_v, "vals": vals * 2})
- df_g = df
- if fill_method is not None:
- df_g = getattr(df.groupby("key"), fill_method)(limit=limit)
- grp = df_g.groupby(df.key)
- expected = grp["vals"].obj / grp["vals"].shift(periods) - 1
- gb = df.groupby("key")
- if frame_or_series is Series:
- gb = gb["vals"]
- else:
- expected = expected.to_frame("vals")
- result = gb.pct_change(
- periods=periods, fill_method=fill_method, limit=limit, freq=freq
- )
- tm.assert_equal(result, expected)
- @pytest.mark.parametrize(
- "func, expected_status",
- [
- ("ffill", ["shrt", "shrt", "lng", np.nan, "shrt", "ntrl", "ntrl"]),
- ("bfill", ["shrt", "lng", "lng", "shrt", "shrt", "ntrl", np.nan]),
- ],
- )
- def test_ffill_bfill_non_unique_multilevel(func, expected_status):
- # GH 19437
- date = pd.to_datetime(
- [
- "2018-01-01",
- "2018-01-01",
- "2018-01-01",
- "2018-01-01",
- "2018-01-02",
- "2018-01-01",
- "2018-01-02",
- ]
- )
- symbol = ["MSFT", "MSFT", "MSFT", "AAPL", "AAPL", "TSLA", "TSLA"]
- status = ["shrt", np.nan, "lng", np.nan, "shrt", "ntrl", np.nan]
- df = DataFrame({"date": date, "symbol": symbol, "status": status})
- df = df.set_index(["date", "symbol"])
- result = getattr(df.groupby("symbol")["status"], func)()
- index = MultiIndex.from_tuples(
- tuples=list(zip(*[date, symbol])), names=["date", "symbol"]
- )
- expected = Series(expected_status, index=index, name="status")
- tm.assert_series_equal(result, expected)
- @pytest.mark.parametrize("func", [np.any, np.all])
- def test_any_all_np_func(func):
- # GH 20653
- df = DataFrame(
- [["foo", True], [np.nan, True], ["foo", True]], columns=["key", "val"]
- )
- exp = Series([True, np.nan, True], name="val")
- res = df.groupby("key")["val"].transform(func)
- tm.assert_series_equal(res, exp)
- def test_groupby_transform_rename():
- # https://github.com/pandas-dev/pandas/issues/23461
- def demean_rename(x):
- result = x - x.mean()
- if isinstance(x, Series):
- return result
- result = result.rename(columns={c: f"{c}_demeaned" for c in result.columns})
- return result
- df = DataFrame({"group": list("ababa"), "value": [1, 1, 1, 2, 2]})
- expected = DataFrame({"value": [-1.0 / 3, -0.5, -1.0 / 3, 0.5, 2.0 / 3]})
- result = df.groupby("group").transform(demean_rename)
- tm.assert_frame_equal(result, expected)
- result_single = df.groupby("group").value.transform(demean_rename)
- tm.assert_series_equal(result_single, expected["value"])
- @pytest.mark.parametrize("func", [min, max, np.min, np.max, "first", "last"])
- def test_groupby_transform_timezone_column(func):
- # GH 24198
- ts = pd.to_datetime("now", utc=True).tz_convert("Asia/Singapore")
- result = DataFrame({"end_time": [ts], "id": [1]})
- result["max_end_time"] = result.groupby("id").end_time.transform(func)
- expected = DataFrame([[ts, 1, ts]], columns=["end_time", "id", "max_end_time"])
- tm.assert_frame_equal(result, expected)
- @pytest.mark.parametrize(
- "func, values",
- [
- ("idxmin", ["1/1/2011"] * 2 + ["1/3/2011"] * 7 + ["1/10/2011"]),
- ("idxmax", ["1/2/2011"] * 2 + ["1/9/2011"] * 7 + ["1/10/2011"]),
- ],
- )
- def test_groupby_transform_with_datetimes(func, values):
- # GH 15306
- dates = date_range("1/1/2011", periods=10, freq="D")
- stocks = DataFrame({"price": np.arange(10.0)}, index=dates)
- stocks["week_id"] = dates.isocalendar().week
- result = stocks.groupby(stocks["week_id"])["price"].transform(func)
- expected = Series(data=pd.to_datetime(values), index=dates, name="price")
- tm.assert_series_equal(result, expected)
- def test_groupby_transform_dtype():
- # GH 22243
- df = DataFrame({"a": [1], "val": [1.35]})
- result = df["val"].transform(lambda x: x.map(lambda y: f"+{y}"))
- expected1 = Series(["+1.35"], name="val", dtype="object")
- tm.assert_series_equal(result, expected1)
- result = df.groupby("a")["val"].transform(lambda x: x.map(lambda y: f"+{y}"))
- tm.assert_series_equal(result, expected1)
- result = df.groupby("a")["val"].transform(lambda x: x.map(lambda y: f"+({y})"))
- expected2 = Series(["+(1.35)"], name="val", dtype="object")
- tm.assert_series_equal(result, expected2)
- df["val"] = df["val"].astype(object)
- result = df.groupby("a")["val"].transform(lambda x: x.map(lambda y: f"+{y}"))
- tm.assert_series_equal(result, expected1)
- @pytest.mark.parametrize("func", ["cumsum", "cumprod", "cummin", "cummax"])
- def test_transform_absent_categories(func):
- # GH 16771
- # cython transforms with more groups than rows
- x_vals = [1]
- x_cats = range(2)
- y = [1]
- df = DataFrame({"x": Categorical(x_vals, x_cats), "y": y})
- result = getattr(df.y.groupby(df.x), func)()
- expected = df.y
- tm.assert_series_equal(result, expected)
- @pytest.mark.parametrize("func", ["ffill", "bfill", "shift"])
- @pytest.mark.parametrize("key, val", [("level", 0), ("by", Series([0]))])
- def test_ffill_not_in_axis(func, key, val):
- # GH 21521
- df = DataFrame([[np.nan]])
- result = getattr(df.groupby(**{key: val}), func)()
- expected = df
- tm.assert_frame_equal(result, expected)
- def test_transform_invalid_name_raises():
- # GH#27486
- df = DataFrame({"a": [0, 1, 1, 2]})
- g = df.groupby(["a", "b", "b", "c"])
- with pytest.raises(ValueError, match="not a valid function name"):
- g.transform("some_arbitrary_name")
- # method exists on the object, but is not a valid transformation/agg
- assert hasattr(g, "aggregate") # make sure the method exists
- with pytest.raises(ValueError, match="not a valid function name"):
- g.transform("aggregate")
- # Test SeriesGroupBy
- g = df["a"].groupby(["a", "b", "b", "c"])
- with pytest.raises(ValueError, match="not a valid function name"):
- g.transform("some_arbitrary_name")
- def test_transform_agg_by_name(request, reduction_func, frame_or_series):
- func = reduction_func
- obj = DataFrame(
- {"a": [0, 0, 0, 1, 1, 1], "b": range(6)},
- index=["A", "B", "C", "D", "E", "F"],
- )
- if frame_or_series is Series:
- obj = obj["a"]
- g = obj.groupby(np.repeat([0, 1], 3))
- if func == "corrwith" and isinstance(obj, Series): # GH#32293
- request.node.add_marker(
- pytest.mark.xfail(reason="TODO: implement SeriesGroupBy.corrwith")
- )
- args = get_groupby_method_args(reduction_func, obj)
- result = g.transform(func, *args)
- # this is the *definition* of a transformation
- tm.assert_index_equal(result.index, obj.index)
- if func not in ("ngroup", "size") and obj.ndim == 2:
- # size/ngroup return a Series, unlike other transforms
- tm.assert_index_equal(result.columns, obj.columns)
- # verify that values were broadcasted across each group
- assert len(set(DataFrame(result).iloc[-3:, -1])) == 1
- def test_transform_lambda_with_datetimetz():
- # GH 27496
- df = DataFrame(
- {
- "time": [
- Timestamp("2010-07-15 03:14:45"),
- Timestamp("2010-11-19 18:47:06"),
- ],
- "timezone": ["Etc/GMT+4", "US/Eastern"],
- }
- )
- result = df.groupby(["timezone"])["time"].transform(
- lambda x: x.dt.tz_localize(x.name)
- )
- expected = Series(
- [
- Timestamp("2010-07-15 03:14:45", tz="Etc/GMT+4"),
- Timestamp("2010-11-19 18:47:06", tz="US/Eastern"),
- ],
- name="time",
- )
- tm.assert_series_equal(result, expected)
- def test_transform_fastpath_raises():
- # GH#29631 case where fastpath defined in groupby.generic _choose_path
- # raises, but slow_path does not
- df = DataFrame({"A": [1, 1, 2, 2], "B": [1, -1, 1, 2]})
- gb = df.groupby("A")
- def func(grp):
- # we want a function such that func(frame) fails but func.apply(frame)
- # works
- if grp.ndim == 2:
- # Ensure that fast_path fails
- raise NotImplementedError("Don't cross the streams")
- return grp * 2
- # Check that the fastpath raises, see _transform_general
- obj = gb._obj_with_exclusions
- gen = gb.grouper.get_iterator(obj, axis=gb.axis)
- fast_path, slow_path = gb._define_paths(func)
- _, group = next(gen)
- with pytest.raises(NotImplementedError, match="Don't cross the streams"):
- fast_path(group)
- result = gb.transform(func)
- expected = DataFrame([2, -2, 2, 4], columns=["B"])
- tm.assert_frame_equal(result, expected)
- def test_transform_lambda_indexing():
- # GH 7883
- df = DataFrame(
- {
- "A": ["foo", "bar", "foo", "bar", "foo", "flux", "foo", "flux"],
- "B": ["one", "one", "two", "three", "two", "six", "five", "three"],
- "C": range(8),
- "D": range(8),
- "E": range(8),
- }
- )
- df = df.set_index(["A", "B"])
- df = df.sort_index()
- result = df.groupby(level="A").transform(lambda x: x.iloc[-1])
- expected = DataFrame(
- {
- "C": [3, 3, 7, 7, 4, 4, 4, 4],
- "D": [3, 3, 7, 7, 4, 4, 4, 4],
- "E": [3, 3, 7, 7, 4, 4, 4, 4],
- },
- index=MultiIndex.from_tuples(
- [
- ("bar", "one"),
- ("bar", "three"),
- ("flux", "six"),
- ("flux", "three"),
- ("foo", "five"),
- ("foo", "one"),
- ("foo", "two"),
- ("foo", "two"),
- ],
- names=["A", "B"],
- ),
- )
- tm.assert_frame_equal(result, expected)
- def test_categorical_and_not_categorical_key(observed):
- # Checks that groupby-transform, when grouping by both a categorical
- # and a non-categorical key, doesn't try to expand the output to include
- # non-observed categories but instead matches the input shape.
- # GH 32494
- df_with_categorical = DataFrame(
- {
- "A": Categorical(["a", "b", "a"], categories=["a", "b", "c"]),
- "B": [1, 2, 3],
- "C": ["a", "b", "a"],
- }
- )
- df_without_categorical = DataFrame(
- {"A": ["a", "b", "a"], "B": [1, 2, 3], "C": ["a", "b", "a"]}
- )
- # DataFrame case
- result = df_with_categorical.groupby(["A", "C"], observed=observed).transform("sum")
- expected = df_without_categorical.groupby(["A", "C"]).transform("sum")
- tm.assert_frame_equal(result, expected)
- expected_explicit = DataFrame({"B": [4, 2, 4]})
- tm.assert_frame_equal(result, expected_explicit)
- # Series case
- result = df_with_categorical.groupby(["A", "C"], observed=observed)["B"].transform(
- "sum"
- )
- expected = df_without_categorical.groupby(["A", "C"])["B"].transform("sum")
- tm.assert_series_equal(result, expected)
- expected_explicit = Series([4, 2, 4], name="B")
- tm.assert_series_equal(result, expected_explicit)
- def test_string_rank_grouping():
- # GH 19354
- df = DataFrame({"A": [1, 1, 2], "B": [1, 2, 3]})
- result = df.groupby("A").transform("rank")
- expected = DataFrame({"B": [1.0, 2.0, 1.0]})
- tm.assert_frame_equal(result, expected)
- def test_transform_cumcount():
- # GH 27472
- df = DataFrame({"a": [0, 0, 0, 1, 1, 1], "b": range(6)})
- grp = df.groupby(np.repeat([0, 1], 3))
- result = grp.cumcount()
- expected = Series([0, 1, 2, 0, 1, 2])
- tm.assert_series_equal(result, expected)
- result = grp.transform("cumcount")
- tm.assert_series_equal(result, expected)
- @pytest.mark.parametrize("keys", [["A1"], ["A1", "A2"]])
- def test_null_group_lambda_self(sort, dropna, keys):
- # GH 17093
- size = 50
- nulls1 = np.random.choice([False, True], size)
- nulls2 = np.random.choice([False, True], size)
- # Whether a group contains a null value or not
- nulls_grouper = nulls1 if len(keys) == 1 else nulls1 | nulls2
- a1 = np.random.randint(0, 5, size=size).astype(float)
- a1[nulls1] = np.nan
- a2 = np.random.randint(0, 5, size=size).astype(float)
- a2[nulls2] = np.nan
- values = np.random.randint(0, 5, size=a1.shape)
- df = DataFrame({"A1": a1, "A2": a2, "B": values})
- expected_values = values
- if dropna and nulls_grouper.any():
- expected_values = expected_values.astype(float)
- expected_values[nulls_grouper] = np.nan
- expected = DataFrame(expected_values, columns=["B"])
- gb = df.groupby(keys, dropna=dropna, sort=sort)
- result = gb[["B"]].transform(lambda x: x)
- tm.assert_frame_equal(result, expected)
- def test_null_group_str_reducer(request, dropna, reduction_func):
- # GH 17093
- if reduction_func == "corrwith":
- msg = "incorrectly raises"
- request.node.add_marker(pytest.mark.xfail(reason=msg))
- index = [1, 2, 3, 4] # test transform preserves non-standard index
- df = DataFrame({"A": [1, 1, np.nan, np.nan], "B": [1, 2, 2, 3]}, index=index)
- gb = df.groupby("A", dropna=dropna)
- args = get_groupby_method_args(reduction_func, df)
- # Manually handle reducers that don't fit the generic pattern
- # Set expected with dropna=False, then replace if necessary
- if reduction_func == "first":
- expected = DataFrame({"B": [1, 1, 2, 2]}, index=index)
- elif reduction_func == "last":
- expected = DataFrame({"B": [2, 2, 3, 3]}, index=index)
- elif reduction_func == "nth":
- expected = DataFrame({"B": [1, 1, 2, 2]}, index=index)
- elif reduction_func == "size":
- expected = Series([2, 2, 2, 2], index=index)
- elif reduction_func == "corrwith":
- expected = DataFrame({"B": [1.0, 1.0, 1.0, 1.0]}, index=index)
- else:
- expected_gb = df.groupby("A", dropna=False)
- buffer = []
- for idx, group in expected_gb:
- res = getattr(group["B"], reduction_func)()
- buffer.append(Series(res, index=group.index))
- expected = concat(buffer).to_frame("B")
- if dropna:
- dtype = object if reduction_func in ("any", "all") else float
- expected = expected.astype(dtype)
- if expected.ndim == 2:
- expected.iloc[[2, 3], 0] = np.nan
- else:
- expected.iloc[[2, 3]] = np.nan
- result = gb.transform(reduction_func, *args)
- tm.assert_equal(result, expected)
- def test_null_group_str_transformer(request, dropna, transformation_func):
- # GH 17093
- df = DataFrame({"A": [1, 1, np.nan], "B": [1, 2, 2]}, index=[1, 2, 3])
- args = get_groupby_method_args(transformation_func, df)
- gb = df.groupby("A", dropna=dropna)
- buffer = []
- for k, (idx, group) in enumerate(gb):
- if transformation_func == "cumcount":
- # DataFrame has no cumcount method
- res = DataFrame({"B": range(len(group))}, index=group.index)
- elif transformation_func == "ngroup":
- res = DataFrame(len(group) * [k], index=group.index, columns=["B"])
- else:
- res = getattr(group[["B"]], transformation_func)(*args)
- buffer.append(res)
- if dropna:
- dtype = object if transformation_func in ("any", "all") else None
- buffer.append(DataFrame([[np.nan]], index=[3], dtype=dtype, columns=["B"]))
- expected = concat(buffer)
- if transformation_func in ("cumcount", "ngroup"):
- # ngroup/cumcount always returns a Series as it counts the groups, not values
- expected = expected["B"].rename(None)
- result = gb.transform(transformation_func, *args)
- tm.assert_equal(result, expected)
- def test_null_group_str_reducer_series(request, dropna, reduction_func):
- # GH 17093
- if reduction_func == "corrwith":
- msg = "corrwith not implemented for SeriesGroupBy"
- request.node.add_marker(pytest.mark.xfail(reason=msg))
- # GH 17093
- index = [1, 2, 3, 4] # test transform preserves non-standard index
- ser = Series([1, 2, 2, 3], index=index)
- gb = ser.groupby([1, 1, np.nan, np.nan], dropna=dropna)
- args = get_groupby_method_args(reduction_func, ser)
- # Manually handle reducers that don't fit the generic pattern
- # Set expected with dropna=False, then replace if necessary
- if reduction_func == "first":
- expected = Series([1, 1, 2, 2], index=index)
- elif reduction_func == "last":
- expected = Series([2, 2, 3, 3], index=index)
- elif reduction_func == "nth":
- expected = Series([1, 1, 2, 2], index=index)
- elif reduction_func == "size":
- expected = Series([2, 2, 2, 2], index=index)
- elif reduction_func == "corrwith":
- expected = Series([1, 1, 2, 2], index=index)
- else:
- expected_gb = ser.groupby([1, 1, np.nan, np.nan], dropna=False)
- buffer = []
- for idx, group in expected_gb:
- res = getattr(group, reduction_func)()
- buffer.append(Series(res, index=group.index))
- expected = concat(buffer)
- if dropna:
- dtype = object if reduction_func in ("any", "all") else float
- expected = expected.astype(dtype)
- expected.iloc[[2, 3]] = np.nan
- result = gb.transform(reduction_func, *args)
- tm.assert_series_equal(result, expected)
- def test_null_group_str_transformer_series(dropna, transformation_func):
- # GH 17093
- ser = Series([1, 2, 2], index=[1, 2, 3])
- args = get_groupby_method_args(transformation_func, ser)
- gb = ser.groupby([1, 1, np.nan], dropna=dropna)
- buffer = []
- for k, (idx, group) in enumerate(gb):
- if transformation_func == "cumcount":
- # Series has no cumcount method
- res = Series(range(len(group)), index=group.index)
- elif transformation_func == "ngroup":
- res = Series(k, index=group.index)
- else:
- res = getattr(group, transformation_func)(*args)
- buffer.append(res)
- if dropna:
- dtype = object if transformation_func in ("any", "all") else None
- buffer.append(Series([np.nan], index=[3], dtype=dtype))
- expected = concat(buffer)
- with tm.assert_produces_warning(None):
- result = gb.transform(transformation_func, *args)
- tm.assert_equal(result, expected)
- @pytest.mark.parametrize(
- "func, expected_values",
- [
- (Series.sort_values, [5, 4, 3, 2, 1]),
- (lambda x: x.head(1), [5.0, np.nan, 3, 2, np.nan]),
- ],
- )
- @pytest.mark.parametrize("keys", [["a1"], ["a1", "a2"]])
- @pytest.mark.parametrize("keys_in_index", [True, False])
- def test_transform_aligns(func, frame_or_series, expected_values, keys, keys_in_index):
- # GH#45648 - transform should align with the input's index
- df = DataFrame({"a1": [1, 1, 3, 2, 2], "b": [5, 4, 3, 2, 1]})
- if "a2" in keys:
- df["a2"] = df["a1"]
- if keys_in_index:
- df = df.set_index(keys, append=True)
- gb = df.groupby(keys)
- if frame_or_series is Series:
- gb = gb["b"]
- result = gb.transform(func)
- expected = DataFrame({"b": expected_values}, index=df.index)
- if frame_or_series is Series:
- expected = expected["b"]
- tm.assert_equal(result, expected)
- @pytest.mark.parametrize("keys", ["A", ["A", "B"]])
- def test_as_index_no_change(keys, df, groupby_func):
- # GH#49834 - as_index should have no impact on DataFrameGroupBy.transform
- if keys == "A":
- # Column B is string dtype; will fail on some ops
- df = df.drop(columns="B")
- args = get_groupby_method_args(groupby_func, df)
- gb_as_index_true = df.groupby(keys, as_index=True)
- gb_as_index_false = df.groupby(keys, as_index=False)
- result = gb_as_index_true.transform(groupby_func, *args)
- expected = gb_as_index_false.transform(groupby_func, *args)
- tm.assert_equal(result, expected)
|