812 lines
24 KiB
Python
812 lines
24 KiB
Python
|
from datetime import datetime
|
||
|
|
from io import StringIO
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
import pandas as pd
|
||
|
|
from pandas import DataFrame, Index, MultiIndex, Series, bdate_range
|
||
|
|
import pandas._testing as tm
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_issues():
|
||
|
|
# GH 5788
|
||
|
|
|
||
|
|
s = """2011.05.16,00:00,1.40893
|
||
|
|
2011.05.16,01:00,1.40760
|
||
|
|
2011.05.16,02:00,1.40750
|
||
|
|
2011.05.16,03:00,1.40649
|
||
|
|
2011.05.17,02:00,1.40893
|
||
|
|
2011.05.17,03:00,1.40760
|
||
|
|
2011.05.17,04:00,1.40750
|
||
|
|
2011.05.17,05:00,1.40649
|
||
|
|
2011.05.18,02:00,1.40893
|
||
|
|
2011.05.18,03:00,1.40760
|
||
|
|
2011.05.18,04:00,1.40750
|
||
|
|
2011.05.18,05:00,1.40649"""
|
||
|
|
|
||
|
|
df = pd.read_csv(
|
||
|
|
StringIO(s),
|
||
|
|
header=None,
|
||
|
|
names=["date", "time", "value"],
|
||
|
|
parse_dates=[["date", "time"]],
|
||
|
|
)
|
||
|
|
df = df.set_index("date_time")
|
||
|
|
|
||
|
|
expected = df.groupby(df.index.date).idxmax()
|
||
|
|
result = df.groupby(df.index.date).apply(lambda x: x.idxmax())
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# GH 5789
|
||
|
|
# don't auto coerce dates
|
||
|
|
df = pd.read_csv(StringIO(s), header=None, names=["date", "time", "value"])
|
||
|
|
exp_idx = pd.Index(
|
||
|
|
["2011.05.16", "2011.05.17", "2011.05.18"], dtype=object, name="date"
|
||
|
|
)
|
||
|
|
expected = Series(["00:00", "02:00", "02:00"], index=exp_idx)
|
||
|
|
result = df.groupby("date").apply(lambda x: x["time"][x["value"].idxmax()])
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_trivial():
|
||
|
|
# GH 20066
|
||
|
|
# trivial apply: ignore input and return a constant dataframe.
|
||
|
|
df = pd.DataFrame(
|
||
|
|
{"key": ["a", "a", "b", "b", "a"], "data": [1.0, 2.0, 3.0, 4.0, 5.0]},
|
||
|
|
columns=["key", "data"],
|
||
|
|
)
|
||
|
|
expected = pd.concat([df.iloc[1:], df.iloc[1:]], axis=1, keys=["float64", "object"])
|
||
|
|
result = df.groupby([str(x) for x in df.dtypes], axis=1).apply(
|
||
|
|
lambda x: df.iloc[1:]
|
||
|
|
)
|
||
|
|
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.xfail(
|
||
|
|
reason="GH#20066; function passed into apply "
|
||
|
|
"returns a DataFrame with the same index "
|
||
|
|
"as the one to create GroupBy object."
|
||
|
|
)
|
||
|
|
def test_apply_trivial_fail():
|
||
|
|
# GH 20066
|
||
|
|
# trivial apply fails if the constant dataframe has the same index
|
||
|
|
# with the one used to create GroupBy object.
|
||
|
|
df = pd.DataFrame(
|
||
|
|
{"key": ["a", "a", "b", "b", "a"], "data": [1.0, 2.0, 3.0, 4.0, 5.0]},
|
||
|
|
columns=["key", "data"],
|
||
|
|
)
|
||
|
|
expected = pd.concat([df, df], axis=1, keys=["float64", "object"])
|
||
|
|
result = df.groupby([str(x) for x in df.dtypes], axis=1).apply(lambda x: df)
|
||
|
|
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_fast_apply():
|
||
|
|
# make sure that fast apply is correctly called
|
||
|
|
# rather than raising any kind of error
|
||
|
|
# otherwise the python path will be callsed
|
||
|
|
# which slows things down
|
||
|
|
N = 1000
|
||
|
|
labels = np.random.randint(0, 2000, size=N)
|
||
|
|
labels2 = np.random.randint(0, 3, size=N)
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"key": labels,
|
||
|
|
"key2": labels2,
|
||
|
|
"value1": np.random.randn(N),
|
||
|
|
"value2": ["foo", "bar", "baz", "qux"] * (N // 4),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
def f(g):
|
||
|
|
return 1
|
||
|
|
|
||
|
|
g = df.groupby(["key", "key2"])
|
||
|
|
|
||
|
|
grouper = g.grouper
|
||
|
|
|
||
|
|
splitter = grouper._get_splitter(g._selected_obj, axis=g.axis)
|
||
|
|
group_keys = grouper._get_group_keys()
|
||
|
|
|
||
|
|
values, mutated = splitter.fast_apply(f, group_keys)
|
||
|
|
|
||
|
|
assert not mutated
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"df, group_names",
|
||
|
|
[
|
||
|
|
(DataFrame({"a": [1, 1, 1, 2, 3], "b": ["a", "a", "a", "b", "c"]}), [1, 2, 3]),
|
||
|
|
(DataFrame({"a": [0, 0, 1, 1], "b": [0, 1, 0, 1]}), [0, 1]),
|
||
|
|
(DataFrame({"a": [1]}), [1]),
|
||
|
|
(DataFrame({"a": [1, 1, 1, 2, 2, 1, 1, 2], "b": range(8)}), [1, 2]),
|
||
|
|
(DataFrame({"a": [1, 2, 3, 1, 2, 3], "two": [4, 5, 6, 7, 8, 9]}), [1, 2, 3]),
|
||
|
|
(
|
||
|
|
DataFrame(
|
||
|
|
{
|
||
|
|
"a": list("aaabbbcccc"),
|
||
|
|
"B": [3, 4, 3, 6, 5, 2, 1, 9, 5, 4],
|
||
|
|
"C": [4, 0, 2, 2, 2, 7, 8, 6, 2, 8],
|
||
|
|
}
|
||
|
|
),
|
||
|
|
["a", "b", "c"],
|
||
|
|
),
|
||
|
|
(DataFrame([[1, 2, 3], [2, 2, 3]], columns=["a", "b", "c"]), [1, 2]),
|
||
|
|
],
|
||
|
|
ids=[
|
||
|
|
"GH2936",
|
||
|
|
"GH7739 & GH10519",
|
||
|
|
"GH10519",
|
||
|
|
"GH2656",
|
||
|
|
"GH12155",
|
||
|
|
"GH20084",
|
||
|
|
"GH21417",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_group_apply_once_per_group(df, group_names):
|
||
|
|
# GH2936, GH7739, GH10519, GH2656, GH12155, GH20084, GH21417
|
||
|
|
|
||
|
|
# This test should ensure that a function is only evaluated
|
||
|
|
# once per group. Previously the function has been evaluated twice
|
||
|
|
# on the first group to check if the Cython index slider is safe to use
|
||
|
|
# This test ensures that the side effect (append to list) is only triggered
|
||
|
|
# once per group
|
||
|
|
|
||
|
|
names = []
|
||
|
|
# cannot parameterize over the functions since they need external
|
||
|
|
# `names` to detect side effects
|
||
|
|
|
||
|
|
def f_copy(group):
|
||
|
|
# this takes the fast apply path
|
||
|
|
names.append(group.name)
|
||
|
|
return group.copy()
|
||
|
|
|
||
|
|
def f_nocopy(group):
|
||
|
|
# this takes the slow apply path
|
||
|
|
names.append(group.name)
|
||
|
|
return group
|
||
|
|
|
||
|
|
def f_scalar(group):
|
||
|
|
# GH7739, GH2656
|
||
|
|
names.append(group.name)
|
||
|
|
return 0
|
||
|
|
|
||
|
|
def f_none(group):
|
||
|
|
# GH10519, GH12155, GH21417
|
||
|
|
names.append(group.name)
|
||
|
|
return None
|
||
|
|
|
||
|
|
def f_constant_df(group):
|
||
|
|
# GH2936, GH20084
|
||
|
|
names.append(group.name)
|
||
|
|
return DataFrame({"a": [1], "b": [1]})
|
||
|
|
|
||
|
|
for func in [f_copy, f_nocopy, f_scalar, f_none, f_constant_df]:
|
||
|
|
del names[:]
|
||
|
|
|
||
|
|
df.groupby("a").apply(func)
|
||
|
|
assert names == group_names
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_with_mixed_dtype():
|
||
|
|
# GH3480, apply with mixed dtype on axis=1 breaks in 0.11
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"foo1": np.random.randn(6),
|
||
|
|
"foo2": ["one", "two", "two", "three", "one", "two"],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
result = df.apply(lambda x: x, axis=1).dtypes
|
||
|
|
expected = df.dtypes
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
# GH 3610 incorrect dtype conversion with as_index=False
|
||
|
|
df = DataFrame({"c1": [1, 2, 6, 6, 8]})
|
||
|
|
df["c2"] = df.c1 / 2.0
|
||
|
|
result1 = df.groupby("c2").mean().reset_index().c2
|
||
|
|
result2 = df.groupby("c2", as_index=False).mean().c2
|
||
|
|
tm.assert_series_equal(result1, result2)
|
||
|
|
|
||
|
|
|
||
|
|
def test_groupby_as_index_apply(df):
|
||
|
|
# GH #4648 and #3417
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"item_id": ["b", "b", "a", "c", "a", "b"],
|
||
|
|
"user_id": [1, 2, 1, 1, 3, 1],
|
||
|
|
"time": range(6),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
g_as = df.groupby("user_id", as_index=True)
|
||
|
|
g_not_as = df.groupby("user_id", as_index=False)
|
||
|
|
|
||
|
|
res_as = g_as.head(2).index
|
||
|
|
res_not_as = g_not_as.head(2).index
|
||
|
|
exp = Index([0, 1, 2, 4])
|
||
|
|
tm.assert_index_equal(res_as, exp)
|
||
|
|
tm.assert_index_equal(res_not_as, exp)
|
||
|
|
|
||
|
|
res_as_apply = g_as.apply(lambda x: x.head(2)).index
|
||
|
|
res_not_as_apply = g_not_as.apply(lambda x: x.head(2)).index
|
||
|
|
|
||
|
|
# apply doesn't maintain the original ordering
|
||
|
|
# changed in GH5610 as the as_index=False returns a MI here
|
||
|
|
exp_not_as_apply = MultiIndex.from_tuples([(0, 0), (0, 2), (1, 1), (2, 4)])
|
||
|
|
tp = [(1, 0), (1, 2), (2, 1), (3, 4)]
|
||
|
|
exp_as_apply = MultiIndex.from_tuples(tp, names=["user_id", None])
|
||
|
|
|
||
|
|
tm.assert_index_equal(res_as_apply, exp_as_apply)
|
||
|
|
tm.assert_index_equal(res_not_as_apply, exp_not_as_apply)
|
||
|
|
|
||
|
|
ind = Index(list("abcde"))
|
||
|
|
df = DataFrame([[1, 2], [2, 3], [1, 4], [1, 5], [2, 6]], index=ind)
|
||
|
|
res = df.groupby(0, as_index=False).apply(lambda x: x).index
|
||
|
|
tm.assert_index_equal(res, ind)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_concat_preserve_names(three_group):
|
||
|
|
grouped = three_group.groupby(["A", "B"])
|
||
|
|
|
||
|
|
def desc(group):
|
||
|
|
result = group.describe()
|
||
|
|
result.index.name = "stat"
|
||
|
|
return result
|
||
|
|
|
||
|
|
def desc2(group):
|
||
|
|
result = group.describe()
|
||
|
|
result.index.name = "stat"
|
||
|
|
result = result[: len(group)]
|
||
|
|
# weirdo
|
||
|
|
return result
|
||
|
|
|
||
|
|
def desc3(group):
|
||
|
|
result = group.describe()
|
||
|
|
|
||
|
|
# names are different
|
||
|
|
result.index.name = f"stat_{len(group):d}"
|
||
|
|
|
||
|
|
result = result[: len(group)]
|
||
|
|
# weirdo
|
||
|
|
return result
|
||
|
|
|
||
|
|
result = grouped.apply(desc)
|
||
|
|
assert result.index.names == ("A", "B", "stat")
|
||
|
|
|
||
|
|
result2 = grouped.apply(desc2)
|
||
|
|
assert result2.index.names == ("A", "B", "stat")
|
||
|
|
|
||
|
|
result3 = grouped.apply(desc3)
|
||
|
|
assert result3.index.names == ("A", "B", None)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_series_to_frame():
|
||
|
|
def f(piece):
|
||
|
|
with np.errstate(invalid="ignore"):
|
||
|
|
logged = np.log(piece)
|
||
|
|
return DataFrame(
|
||
|
|
{"value": piece, "demeaned": piece - piece.mean(), "logged": logged}
|
||
|
|
)
|
||
|
|
|
||
|
|
dr = bdate_range("1/1/2000", periods=100)
|
||
|
|
ts = Series(np.random.randn(100), index=dr)
|
||
|
|
|
||
|
|
grouped = ts.groupby(lambda x: x.month)
|
||
|
|
result = grouped.apply(f)
|
||
|
|
|
||
|
|
assert isinstance(result, DataFrame)
|
||
|
|
tm.assert_index_equal(result.index, ts.index)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_series_yield_constant(df):
|
||
|
|
result = df.groupby(["A", "B"])["C"].apply(len)
|
||
|
|
assert result.index.names[:2] == ("A", "B")
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_frame_yield_constant(df):
|
||
|
|
# GH13568
|
||
|
|
result = df.groupby(["A", "B"]).apply(len)
|
||
|
|
assert isinstance(result, Series)
|
||
|
|
assert result.name is None
|
||
|
|
|
||
|
|
result = df.groupby(["A", "B"])[["C", "D"]].apply(len)
|
||
|
|
assert isinstance(result, Series)
|
||
|
|
assert result.name is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_frame_to_series(df):
|
||
|
|
grouped = df.groupby(["A", "B"])
|
||
|
|
result = grouped.apply(len)
|
||
|
|
expected = grouped.count()["C"]
|
||
|
|
tm.assert_index_equal(result.index, expected.index)
|
||
|
|
tm.assert_numpy_array_equal(result.values, expected.values)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_frame_concat_series():
|
||
|
|
def trans(group):
|
||
|
|
return group.groupby("B")["C"].sum().sort_values()[:2]
|
||
|
|
|
||
|
|
def trans2(group):
|
||
|
|
grouped = group.groupby(df.reindex(group.index)["B"])
|
||
|
|
return grouped.sum().sort_values()[:2]
|
||
|
|
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"A": np.random.randint(0, 5, 1000),
|
||
|
|
"B": np.random.randint(0, 5, 1000),
|
||
|
|
"C": np.random.randn(1000),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
result = df.groupby("A").apply(trans)
|
||
|
|
exp = df.groupby("A")["C"].apply(trans2)
|
||
|
|
tm.assert_series_equal(result, exp, check_names=False)
|
||
|
|
assert result.name == "C"
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_transform(ts):
|
||
|
|
grouped = ts.groupby(lambda x: x.month)
|
||
|
|
result = grouped.apply(lambda x: x * 2)
|
||
|
|
expected = grouped.transform(lambda x: x * 2)
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_multikey_corner(tsframe):
|
||
|
|
grouped = tsframe.groupby([lambda x: x.year, lambda x: x.month])
|
||
|
|
|
||
|
|
def f(group):
|
||
|
|
return group.sort_values("A")[-5:]
|
||
|
|
|
||
|
|
result = grouped.apply(f)
|
||
|
|
for key, group in grouped:
|
||
|
|
tm.assert_frame_equal(result.loc[key], f(group))
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_chunk_view():
|
||
|
|
# Low level tinkering could be unsafe, make sure not
|
||
|
|
df = DataFrame({"key": [1, 1, 1, 2, 2, 2, 3, 3, 3], "value": range(9)})
|
||
|
|
|
||
|
|
result = df.groupby("key", group_keys=False).apply(lambda x: x[:2])
|
||
|
|
expected = df.take([0, 1, 3, 4, 6, 7])
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_no_name_column_conflict():
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"name": [1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
|
||
|
|
"name2": [0, 0, 0, 1, 1, 1, 0, 0, 1, 1],
|
||
|
|
"value": range(9, -1, -1),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
# it works! #2605
|
||
|
|
grouped = df.groupby(["name", "name2"])
|
||
|
|
grouped.apply(lambda x: x.sort_values("value", inplace=True))
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_typecast_fail():
|
||
|
|
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),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
def f(group):
|
||
|
|
v = group["v"]
|
||
|
|
group["v2"] = (v - v.min()) / (v.max() - v.min())
|
||
|
|
return group
|
||
|
|
|
||
|
|
result = df.groupby("d").apply(f)
|
||
|
|
|
||
|
|
expected = df.copy()
|
||
|
|
expected["v2"] = np.tile([0.0, 0.5, 1], 2)
|
||
|
|
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_multiindex_fail():
|
||
|
|
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):
|
||
|
|
v = group["v"]
|
||
|
|
group["v2"] = (v - v.min()) / (v.max() - v.min())
|
||
|
|
return group
|
||
|
|
|
||
|
|
result = df.groupby("d").apply(f)
|
||
|
|
|
||
|
|
expected = df.copy()
|
||
|
|
expected["v2"] = np.tile([0.0, 0.5, 1], 2)
|
||
|
|
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_corner(tsframe):
|
||
|
|
result = tsframe.groupby(lambda x: x.year).apply(lambda x: x * 2)
|
||
|
|
expected = tsframe * 2
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_without_copy():
|
||
|
|
# GH 5545
|
||
|
|
# returning a non-copy in an applied function fails
|
||
|
|
|
||
|
|
data = DataFrame(
|
||
|
|
{
|
||
|
|
"id_field": [100, 100, 200, 300],
|
||
|
|
"category": ["a", "b", "c", "c"],
|
||
|
|
"value": [1, 2, 3, 4],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
def filt1(x):
|
||
|
|
if x.shape[0] == 1:
|
||
|
|
return x.copy()
|
||
|
|
else:
|
||
|
|
return x[x.category == "c"]
|
||
|
|
|
||
|
|
def filt2(x):
|
||
|
|
if x.shape[0] == 1:
|
||
|
|
return x
|
||
|
|
else:
|
||
|
|
return x[x.category == "c"]
|
||
|
|
|
||
|
|
expected = data.groupby("id_field").apply(filt1)
|
||
|
|
result = data.groupby("id_field").apply(filt2)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_corner_cases():
|
||
|
|
# #535, can't use sliding iterator
|
||
|
|
|
||
|
|
N = 1000
|
||
|
|
labels = np.random.randint(0, 100, size=N)
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"key": labels,
|
||
|
|
"value1": np.random.randn(N),
|
||
|
|
"value2": ["foo", "bar", "baz", "qux"] * (N // 4),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
grouped = df.groupby("key")
|
||
|
|
|
||
|
|
def f(g):
|
||
|
|
g["value3"] = g["value1"] * 2
|
||
|
|
return g
|
||
|
|
|
||
|
|
result = grouped.apply(f)
|
||
|
|
assert "value3" in result
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_numeric_coercion_when_datetime():
|
||
|
|
# In the past, group-by/apply operations have been over-eager
|
||
|
|
# in converting dtypes to numeric, in the presence of datetime
|
||
|
|
# columns. Various GH issues were filed, the reproductions
|
||
|
|
# for which are here.
|
||
|
|
|
||
|
|
# GH 15670
|
||
|
|
df = pd.DataFrame(
|
||
|
|
{"Number": [1, 2], "Date": ["2017-03-02"] * 2, "Str": ["foo", "inf"]}
|
||
|
|
)
|
||
|
|
expected = df.groupby(["Number"]).apply(lambda x: x.iloc[0])
|
||
|
|
df.Date = pd.to_datetime(df.Date)
|
||
|
|
result = df.groupby(["Number"]).apply(lambda x: x.iloc[0])
|
||
|
|
tm.assert_series_equal(result["Str"], expected["Str"])
|
||
|
|
|
||
|
|
# GH 15421
|
||
|
|
df = pd.DataFrame(
|
||
|
|
{"A": [10, 20, 30], "B": ["foo", "3", "4"], "T": [pd.Timestamp("12:31:22")] * 3}
|
||
|
|
)
|
||
|
|
|
||
|
|
def get_B(g):
|
||
|
|
return g.iloc[0][["B"]]
|
||
|
|
|
||
|
|
result = df.groupby("A").apply(get_B)["B"]
|
||
|
|
expected = df.B
|
||
|
|
expected.index = df.A
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
# GH 14423
|
||
|
|
def predictions(tool):
|
||
|
|
out = pd.Series(index=["p1", "p2", "useTime"], dtype=object)
|
||
|
|
if "step1" in list(tool.State):
|
||
|
|
out["p1"] = str(tool[tool.State == "step1"].Machine.values[0])
|
||
|
|
if "step2" in list(tool.State):
|
||
|
|
out["p2"] = str(tool[tool.State == "step2"].Machine.values[0])
|
||
|
|
out["useTime"] = str(tool[tool.State == "step2"].oTime.values[0])
|
||
|
|
return out
|
||
|
|
|
||
|
|
df1 = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"Key": ["B", "B", "A", "A"],
|
||
|
|
"State": ["step1", "step2", "step1", "step2"],
|
||
|
|
"oTime": ["", "2016-09-19 05:24:33", "", "2016-09-19 23:59:04"],
|
||
|
|
"Machine": ["23", "36L", "36R", "36R"],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
df2 = df1.copy()
|
||
|
|
df2.oTime = pd.to_datetime(df2.oTime)
|
||
|
|
expected = df1.groupby("Key").apply(predictions).p1
|
||
|
|
result = df2.groupby("Key").apply(predictions).p1
|
||
|
|
tm.assert_series_equal(expected, result)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_aggregating_timedelta_and_datetime():
|
||
|
|
# Regression test for GH 15562
|
||
|
|
# The following groupby caused ValueErrors and IndexErrors pre 0.20.0
|
||
|
|
|
||
|
|
df = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"clientid": ["A", "B", "C"],
|
||
|
|
"datetime": [np.datetime64("2017-02-01 00:00:00")] * 3,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
df["time_delta_zero"] = df.datetime - df.datetime
|
||
|
|
result = df.groupby("clientid").apply(
|
||
|
|
lambda ddf: pd.Series(
|
||
|
|
dict(clientid_age=ddf.time_delta_zero.min(), date=ddf.datetime.min())
|
||
|
|
)
|
||
|
|
)
|
||
|
|
expected = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"clientid": ["A", "B", "C"],
|
||
|
|
"clientid_age": [np.timedelta64(0, "D")] * 3,
|
||
|
|
"date": [np.datetime64("2017-02-01 00:00:00")] * 3,
|
||
|
|
}
|
||
|
|
).set_index("clientid")
|
||
|
|
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_time_field_bug():
|
||
|
|
# Test a fix for the following error related to GH issue 11324 When
|
||
|
|
# non-key fields in a group-by dataframe contained time-based fields
|
||
|
|
# that were not returned by the apply function, an exception would be
|
||
|
|
# raised.
|
||
|
|
|
||
|
|
df = pd.DataFrame({"a": 1, "b": [datetime.now() for nn in range(10)]})
|
||
|
|
|
||
|
|
def func_with_no_date(batch):
|
||
|
|
return pd.Series({"c": 2})
|
||
|
|
|
||
|
|
def func_with_date(batch):
|
||
|
|
return pd.Series({"b": datetime(2015, 1, 1), "c": 2})
|
||
|
|
|
||
|
|
dfg_no_conversion = df.groupby(by=["a"]).apply(func_with_no_date)
|
||
|
|
dfg_no_conversion_expected = pd.DataFrame({"c": 2}, index=[1])
|
||
|
|
dfg_no_conversion_expected.index.name = "a"
|
||
|
|
|
||
|
|
dfg_conversion = df.groupby(by=["a"]).apply(func_with_date)
|
||
|
|
dfg_conversion_expected = pd.DataFrame(
|
||
|
|
{"b": datetime(2015, 1, 1), "c": 2}, index=[1]
|
||
|
|
)
|
||
|
|
dfg_conversion_expected.index.name = "a"
|
||
|
|
|
||
|
|
tm.assert_frame_equal(dfg_no_conversion, dfg_no_conversion_expected)
|
||
|
|
tm.assert_frame_equal(dfg_conversion, dfg_conversion_expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_gb_apply_list_of_unequal_len_arrays():
|
||
|
|
|
||
|
|
# GH1738
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"group1": ["a", "a", "a", "b", "b", "b", "a", "a", "a", "b", "b", "b"],
|
||
|
|
"group2": ["c", "c", "d", "d", "d", "e", "c", "c", "d", "d", "d", "e"],
|
||
|
|
"weight": [1.1, 2, 3, 4, 5, 6, 2, 4, 6, 8, 1, 2],
|
||
|
|
"value": [7.1, 8, 9, 10, 11, 12, 8, 7, 6, 5, 4, 3],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
df = df.set_index(["group1", "group2"])
|
||
|
|
df_grouped = df.groupby(level=["group1", "group2"], sort=True)
|
||
|
|
|
||
|
|
def noddy(value, weight):
|
||
|
|
out = np.array(value * weight).repeat(3)
|
||
|
|
return out
|
||
|
|
|
||
|
|
# the kernel function returns arrays of unequal length
|
||
|
|
# pandas sniffs the first one, sees it's an array and not
|
||
|
|
# a list, and assumed the rest are of equal length
|
||
|
|
# and so tries a vstack
|
||
|
|
|
||
|
|
# don't die
|
||
|
|
df_grouped.apply(lambda x: noddy(x.value, x.weight))
|
||
|
|
|
||
|
|
|
||
|
|
def test_groupby_apply_all_none():
|
||
|
|
# Tests to make sure no errors if apply function returns all None
|
||
|
|
# values. Issue 9684.
|
||
|
|
test_df = DataFrame({"groups": [0, 0, 1, 1], "random_vars": [8, 7, 4, 5]})
|
||
|
|
|
||
|
|
def test_func(x):
|
||
|
|
pass
|
||
|
|
|
||
|
|
result = test_df.groupby("groups").apply(test_func)
|
||
|
|
expected = DataFrame()
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_groupby_apply_none_first():
|
||
|
|
# GH 12824. Tests if apply returns None first.
|
||
|
|
test_df1 = DataFrame({"groups": [1, 1, 1, 2], "vars": [0, 1, 2, 3]})
|
||
|
|
test_df2 = DataFrame({"groups": [1, 2, 2, 2], "vars": [0, 1, 2, 3]})
|
||
|
|
|
||
|
|
def test_func(x):
|
||
|
|
if x.shape[0] < 2:
|
||
|
|
return None
|
||
|
|
return x.iloc[[0, -1]]
|
||
|
|
|
||
|
|
result1 = test_df1.groupby("groups").apply(test_func)
|
||
|
|
result2 = test_df2.groupby("groups").apply(test_func)
|
||
|
|
index1 = MultiIndex.from_arrays([[1, 1], [0, 2]], names=["groups", None])
|
||
|
|
index2 = MultiIndex.from_arrays([[2, 2], [1, 3]], names=["groups", None])
|
||
|
|
expected1 = DataFrame({"groups": [1, 1], "vars": [0, 2]}, index=index1)
|
||
|
|
expected2 = DataFrame({"groups": [2, 2], "vars": [1, 3]}, index=index2)
|
||
|
|
tm.assert_frame_equal(result1, expected1)
|
||
|
|
tm.assert_frame_equal(result2, expected2)
|
||
|
|
|
||
|
|
|
||
|
|
def test_groupby_apply_return_empty_chunk():
|
||
|
|
# GH 22221: apply filter which returns some empty groups
|
||
|
|
df = pd.DataFrame(dict(value=[0, 1], group=["filled", "empty"]))
|
||
|
|
groups = df.groupby("group")
|
||
|
|
result = groups.apply(lambda group: group[group.value != 1]["value"])
|
||
|
|
expected = pd.Series(
|
||
|
|
[0],
|
||
|
|
name="value",
|
||
|
|
index=MultiIndex.from_product(
|
||
|
|
[["empty", "filled"], [0]], names=["group", None]
|
||
|
|
).drop("empty"),
|
||
|
|
)
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_with_mixed_types():
|
||
|
|
# gh-20949
|
||
|
|
df = pd.DataFrame({"A": "a a b".split(), "B": [1, 2, 3], "C": [4, 6, 5]})
|
||
|
|
g = df.groupby("A")
|
||
|
|
|
||
|
|
result = g.transform(lambda x: x / x.sum())
|
||
|
|
expected = pd.DataFrame({"B": [1 / 3.0, 2 / 3.0, 1], "C": [0.4, 0.6, 1.0]})
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result = g.apply(lambda x: x / x.sum())
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_func_returns_object():
|
||
|
|
# GH 28652
|
||
|
|
df = DataFrame({"a": [1, 2]}, index=pd.Int64Index([1, 2]))
|
||
|
|
result = df.groupby("a").apply(lambda g: g.index)
|
||
|
|
expected = Series(
|
||
|
|
[pd.Int64Index([1]), pd.Int64Index([2])], index=pd.Int64Index([1, 2], name="a")
|
||
|
|
)
|
||
|
|
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"group_column_dtlike",
|
||
|
|
[datetime.today(), datetime.today().date(), datetime.today().time()],
|
||
|
|
)
|
||
|
|
def test_apply_datetime_issue(group_column_dtlike):
|
||
|
|
# GH-28247
|
||
|
|
# groupby-apply throws an error if one of the columns in the DataFrame
|
||
|
|
# is a datetime object and the column labels are different from
|
||
|
|
# standard int values in range(len(num_columns))
|
||
|
|
|
||
|
|
df = pd.DataFrame({"a": ["foo"], "b": [group_column_dtlike]})
|
||
|
|
result = df.groupby("a").apply(lambda x: pd.Series(["spam"], index=[42]))
|
||
|
|
|
||
|
|
expected = pd.DataFrame(
|
||
|
|
["spam"], Index(["foo"], dtype="object", name="a"), columns=[42]
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_apply_series_return_dataframe_groups():
|
||
|
|
# GH 10078
|
||
|
|
tdf = DataFrame(
|
||
|
|
{
|
||
|
|
"day": {
|
||
|
|
0: pd.Timestamp("2015-02-24 00:00:00"),
|
||
|
|
1: pd.Timestamp("2015-02-24 00:00:00"),
|
||
|
|
2: pd.Timestamp("2015-02-24 00:00:00"),
|
||
|
|
3: pd.Timestamp("2015-02-24 00:00:00"),
|
||
|
|
4: pd.Timestamp("2015-02-24 00:00:00"),
|
||
|
|
},
|
||
|
|
"userAgent": {
|
||
|
|
0: "some UA string",
|
||
|
|
1: "some UA string",
|
||
|
|
2: "some UA string",
|
||
|
|
3: "another UA string",
|
||
|
|
4: "some UA string",
|
||
|
|
},
|
||
|
|
"userId": {
|
||
|
|
0: "17661101",
|
||
|
|
1: "17661101",
|
||
|
|
2: "17661101",
|
||
|
|
3: "17661101",
|
||
|
|
4: "17661101",
|
||
|
|
},
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
def most_common_values(df):
|
||
|
|
return Series({c: s.value_counts().index[0] for c, s in df.iteritems()})
|
||
|
|
|
||
|
|
result = tdf.groupby("day").apply(most_common_values)["userId"]
|
||
|
|
expected = pd.Series(
|
||
|
|
["17661101"], index=pd.DatetimeIndex(["2015-02-24"], name="day"), name="userId"
|
||
|
|
)
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("category", [False, True])
|
||
|
|
def test_apply_multi_level_name(category):
|
||
|
|
# https://github.com/pandas-dev/pandas/issues/31068
|
||
|
|
b = [1, 2] * 5
|
||
|
|
if category:
|
||
|
|
b = pd.Categorical(b, categories=[1, 2, 3])
|
||
|
|
df = pd.DataFrame(
|
||
|
|
{"A": np.arange(10), "B": b, "C": list(range(10)), "D": list(range(10))}
|
||
|
|
).set_index(["A", "B"])
|
||
|
|
result = df.groupby("B").apply(lambda x: x.sum())
|
||
|
|
expected = pd.DataFrame(
|
||
|
|
{"C": [20, 25], "D": [20, 25]}, index=pd.Index([1, 2], name="B")
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
assert df.index.names == ["A", "B"]
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"index",
|
||
|
|
[
|
||
|
|
pd.CategoricalIndex(list("abc")),
|
||
|
|
pd.interval_range(0, 3),
|
||
|
|
pd.period_range("2020", periods=3, freq="D"),
|
||
|
|
pd.MultiIndex.from_tuples([("a", 0), ("a", 1), ("b", 0)]),
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_apply_index_has_complex_internals(index):
|
||
|
|
# GH 31248
|
||
|
|
df = DataFrame({"group": [1, 1, 2], "value": [0, 1, 0]}, index=index)
|
||
|
|
result = df.groupby("group").apply(lambda x: x)
|
||
|
|
tm.assert_frame_equal(result, df)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"function, expected_values",
|
||
|
|
[
|
||
|
|
(lambda x: x.index.to_list(), [[0, 1], [2, 3]]),
|
||
|
|
(lambda x: set(x.index.to_list()), [{0, 1}, {2, 3}]),
|
||
|
|
(lambda x: tuple(x.index.to_list()), [(0, 1), (2, 3)]),
|
||
|
|
(
|
||
|
|
lambda x: {n: i for (n, i) in enumerate(x.index.to_list())},
|
||
|
|
[{0: 0, 1: 1}, {0: 2, 1: 3}],
|
||
|
|
),
|
||
|
|
(
|
||
|
|
lambda x: [{n: i} for (n, i) in enumerate(x.index.to_list())],
|
||
|
|
[[{0: 0}, {1: 1}], [{0: 2}, {1: 3}]],
|
||
|
|
),
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_apply_function_returns_non_pandas_non_scalar(function, expected_values):
|
||
|
|
# GH 31441
|
||
|
|
df = pd.DataFrame(["A", "A", "B", "B"], columns=["groups"])
|
||
|
|
result = df.groupby("groups").apply(function)
|
||
|
|
expected = pd.Series(expected_values, index=pd.Index(["A", "B"], name="groups"))
|
||
|
|
tm.assert_series_equal(result, expected)
|