使用更快的方法执行pandas groupby操作

12

我有一个数据集,有姓名(person_name)、天数和颜色(shirt_color)三列。

每个人在特定的一天穿着某种颜色的衬衫。天数可以是任意的。

例如,输入:

name  day  color
----------------
John   1   White
John   2   White
John   3   Blue
John   4   Blue
John   5   White
Tom    2   White
Tom    3   Blue
Tom    4   Blue
Tom    5   Black
Jerry  1   Black
Jerry  2   Black
Jerry  4   Black
Jerry  5   White

我需要找出每个人最常用的颜色。

例如,结果应如下:

name    color
-------------
Jerry   Black
John    White
Tom     Blue

我正在执行以下操作来获得结果,这个方法能够正常工作但速度较慢:

most_frquent_list = [[name, group.color.mode()[0]] 
                        for name, group in data.groupby('name')]
most_frquent_df = pd.DataFrame(most_frquent_list, columns=['name', 'color'])

假设我有一个包含500万个独特名称的数据集。执行上述操作的最佳/最快方法是什么?


我希望有人能对所有这些提交进行基准测试。我现在可以做,但这里已经很晚了。 - André C. Andersen
@AndréC.Andersen 我会为每个解决方案添加基准测试的注释。 - DYZ
1.91毫秒±2.35微秒每个循环(平均值±7次运行的标准差,每个循环1000次) - DYZ
7个回答

17

Numpy的numpy.add.atpandas.factorize

这旨在追求速度。然而,我也尽力将其组织得易于阅读。

i, r = pd.factorize(df.name)
j, c = pd.factorize(df.color)
n, m = len(r), len(c)

b = np.zeros((n, m), dtype=np.int64)

np.add.at(b, (i, j), 1)
pd.Series(c[b.argmax(1)], r)

John     White
Tom       Blue
Jerry    Black
dtype: object

groupbysizeidxmax

df.groupby(['name', 'color']).size().unstack().idxmax(1)

name
Jerry    Black
John     White
Tom       Blue
dtype: object

name
Jerry    Black
John     White
Tom       Blue
Name: color, dtype: object

计数器

¯\_(ツ)_/¯

from collections import Counter

df.groupby('name').color.apply(lambda c: Counter(c).most_common(1)[0][0])

name
Jerry    Black
John     White
Tom       Blue
Name: color, dtype: object

6
1st: 362微秒±1.47微秒每次循环(7次运行的平均值和标准偏差,每次1000次循环) - DYZ
4
第二个测试结果:平均每次循环1.51毫秒,标准差为4.67微秒,共进行7次测试,每次测试循环1000次。 - DYZ
4
第三次测试的结果为834微秒,标准差为2.66微秒,共进行了7次试验,每次试验循环执行了1000次。 - DYZ

6

更新

要打败这个速度(在样本数据框架上比任何提出的pandas解决方案快约10倍,比提出的numpy解决方案快1.5倍)肯定很难。 要点是远离pandas,使用itertools.groupby来处理非数字数据时会更好。

from itertools import groupby
from collections import Counter

pd.Series({x: Counter(z[-1] for z in y).most_common(1)[0][0] for x,y 
          in groupby(sorted(df.values.tolist()), 
                            key=lambda x: x[0])})
# Jerry    Black
# John     White
# Tom       Blue

旧答案

这里有另一种方法。实际上比原来的方法慢,但是我会保留它:

data.groupby('name')['color']\
    .apply(pd.Series.value_counts)\
    .unstack().idxmax(axis=1)
# name
# Jerry    Black
# John     White
# Tom       Blue

哈!我刚刚做到了。我会删除的。 - piRSquared
@piRSquared,拜托了,保留它!让原帖作者自己决定。 - DYZ
@piRSquared,你的“Counter”仍然比较慢,因为使用了“apply”。这里的重点不是要搞乱pandas。 - DYZ
我认为collections的解决方案很好,但是将其称为比pandas/numpy快10倍是误导性的。即使是只有几百行的数据框,在piRSquared的因子化解决方案下也能轻松击败它,并且在样本数据框上的时间测试并不意味着太多。 - user3483203
@user3483203 同意。我添加了一条注释,说明10倍加速只在示例数据帧上看到。 - DYZ
显示剩余2条评论

5
pd.Series.mode的解决方案
df.groupby('name').color.apply(pd.Series.mode).reset_index(level=1,drop=True)
Out[281]: 
name
Jerry    Black
John     White
Tom       Blue
Name: color, dtype: object

抱歉,我误解了问题并进行了修正。 - BENY
每个循环1.66毫秒±3.48微秒(平均值±7次运行的标准偏差,每个循环1000次)。 - DYZ

2
如何使用transform(max)进行两次分组?
df = df.groupby(["name", "color"], as_index=False, sort=False).count()
idx = df.groupby("name", sort=False).transform(max)["day"] == df["day"]
df = df[idx][["name", "color"]].reset_index(drop=True)

输出:

    name  color
0   John  White
1    Tom   Blue
2  Jerry  Black

每次循环平均需要12.2毫秒,标准差为48.4微秒(经过7次运行和100次循环的平均值)。 - DYZ
1
谢谢。根据我所看到的,可能不是很好。对于时间测试的反馈,如果您多次循环相同的小数据集,它可能与在大数据集上执行一次不可比较。许多解决方案具有高启动成本,但一旦开始处理,性能就会良好。多次循环小数据集意味着您可能只是在测量启动成本,这应该只是一次成本。我建议您增加正在测试的数据集的大小,直到运行一个循环需要几秒钟。 - André C. Andersen

1

类似于 @piRSquared 的 pd.factorizenp.add.at 答案。

我们使用编码对列中的字符串进行编码。

i, r = pd.factorize(df.name)
j, c = pd.factorize(df.color)
n, m = len(r), len(c)
b = np.zeros((n, m), dtype=np.int64)

但是,相反地,不要这样做:

np.add.at(b, (i, j), 1)
max_columns_after_add_at = b.argmax(1)

我们使用一个JIT函数来获取max_columns_after_add_at,以便在同一循环中执行添加和查找最大值的操作:
@nb.jit(nopython=True, cache=True)
def add_at(x, rows, cols, val):
    max_vals = np.zeros((x.shape[0], ), np.int64)
    max_inds = np.zeros((x.shape[0], ), np.int64)
    for i in range(len(rows)):
        r = rows[i]
        c = cols[i]
        x[r, c]+=1
        if(x[r, c] > max_vals[r]):
            max_vals[r] = x[r, c]
            max_inds[r] = c
    return max_inds

最终获得数据框。
ans = pd.Series(c[max_columns_after_add_at], r)

所以,区别在于我们如何在 np.add.at() 之后执行 argmax(axis=1)。
时间分析
import numpy as np
import numba as nb
m = 100000
n = 100000
rows = np.random.randint(low = 0, high = m, size=10000)
cols = np.random.randint(low = 0, high = n, size=10000)

所以这个:

%%time
x = np.zeros((m,n))
np.add.at(x, (rows, cols), 1)
maxs = x.argmax(1)

需要翻译的内容:

CPU时间:用户12.4秒,系统38秒,总共50.4秒 墙时:50.5秒

还有这个

%%time
x = np.zeros((m,n))
maxs2 = add_at(x, rows, cols, 1)

给出

CPU时间:用户108毫秒,系统39.4秒,总共39.5秒 墙上时间:38.4秒


1
大多数其他答案中讨论的测试结果都存在偏差,因为使用了一个微不足道的小测试 DataFrame 进行测量。Pandas 有一些固定但通常可以忽略不计的设置时间,但在处理这个微小数据集时它将显得非常重要。
在较大的数据集上,最快的方法是使用 pd.Series.mode() 和 agg():
df.groupby('name')['color'].agg(pd.Series.mode)

测试台:
arr = np.array([
    ('John',   1,   'White'),
    ('John',   2,  'White'),
    ('John',   3,   'Blue'),
    ('John',   4,   'Blue'),
    ('John',   5,   'White'),
    ('Tom',    2,   'White'),
    ('Tom',    3,   'Blue'),
    ('Tom',    4,   'Blue'),
    ('Tom',    5,   'Black'),
    ('Jerry',  1,   'Black'),
    ('Jerry',  2,   'Black'),
    ('Jerry',  4,   'Black'),
    ('Jerry',  5,   'White')],
    dtype=[('name', 'O'), ('day', 'i8'), ('color', 'O')])

from timeit import Timer
from itertools import groupby
from collections import Counter

df = pd.DataFrame.from_records(arr).sample(100_000, replace=True)

def factorize():
    i, r = pd.factorize(df.name)
    j, c = pd.factorize(df.color)
    n, m = len(r), len(c)

    b = np.zeros((n, m), dtype=np.int64)

    np.add.at(b, (i, j), 1)
    return pd.Series(c[b.argmax(1)], r)

t_factorize = Timer(lambda: factorize())
t_idxmax = Timer(lambda: df.groupby(['name', 'color']).size().unstack().idxmax(1))
t_aggmode = Timer(lambda: df.groupby('name')['color'].agg(pd.Series.mode))
t_applymode = Timer(lambda: df.groupby('name').color.apply(pd.Series.mode).reset_index(level=1,drop=True))
t_aggcounter = Timer(lambda: df.groupby('name')['color'].agg(lambda c: Counter(c).most_common(1)[0][0]))
t_applycounter = Timer(lambda: df.groupby('name').color.apply(lambda c: Counter(c).most_common(1)[0][0]))
t_itertools = Timer(lambda: pd.Series(
    {x: Counter(z[-1] for z in y).most_common(1)[0][0] for x,y
      in groupby(sorted(df.values.tolist()), key=lambda x: x[0])}))

n = 100
[print(r) for r in (
    f"{t_factorize.timeit(number=n)=}",
    f"{t_idxmax.timeit(number=n)=}",
    f"{t_aggmode.timeit(number=n)=}",
    f"{t_applymode.timeit(number=n)=}",
    f"{t_applycounter.timeit(number=n)=}",
    f"{t_aggcounter.timeit(number=n)=}",
    f"{t_itertools.timeit(number=n)=}",
)]

t_factorize.timeit(number=n)=1.325189442
t_idxmax.timeit(number=n)=1.0613339019999999
t_aggmode.timeit(number=n)=1.0495010750000002
t_applymode.timeit(number=n)=1.2837302849999999
t_applycounter.timeit(number=n)=1.9432825890000007
t_aggcounter.timeit(number=n)=1.8283823839999993
t_itertools.timeit(number=n)=7.0855046380000015

0
对于想要将上述表格转换为数据框并尝试发布的答案的人,可以使用以下代码片段。将上面的表格复制粘贴到笔记本电脑单元格中,如下所示,请确保删除连字符。
l = """name  day  color
John   1   White
John   2   White
John   3   Blue
John   4   Blue
John   5   White
Tom    2   White
Tom    3   Blue
Tom    4   Blue
Tom    5   Black
Jerry  1   Black
Jerry  2   Black
Jerry  4   Black
Jerry  5   White""".split('\n')

现在我们需要将这个列表转换成元组的列表。
df = pd.DataFrame([tuple(i.split()) for i in l])
headers = df.iloc[0]
new_df  = pd.DataFrame(df.values[1:], columns=headers)

使用 new_df 现在,您可以参考 @piRSquared 上面的答案。

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接