选择每组的最大行数 - 熊猫性能问题

Posted 2023-03-11

【中文标题】选择每组的最大行数 - 熊猫性能问题

【英文标题】：Select the max row per group - pandas performance issue

【发布时间】：2018-10-27 02:37:14

【问题描述】：

我选择每组最多一行，我使用groupby/agg 返回索引值并使用loc 选择行。

例如，按"Id"分组，然后选择"delta"值最高的行：

selected_idx = df.groupby("Id").apply(lambda df: df.delta.argmax())
selected_rows = df.loc[selected_idx, :]

但是，这种方式太慢了。实际上，当我在 1300 万行上使用此查询时，我的 i7/16G RAM 笔记本电脑会挂起。

我有两个问题要请教专家：

如何使这个查询在 pandas 中快速运行？我做错了什么？ 为什么这个手术这么贵？

---

[更新] 非常感谢@unutbu 的分析！ sort_drop 是！在我的 i7/32GRAM 机器上，groupby+idxmax 挂了将近 14 个小时（从不返回任何东西）但是 sort_drop 不到一分钟就处理了！

我仍然需要看看 pandas 是如何实现每个方法的，但问题现在已经解决了！我喜欢 ***。

【问题讨论】：

相关：How do I find: Is the first non-NaN value in each column the maximum for that column in a DataFrame?。也很高兴看到这些答案的表现如何比较和扩展。

【参考方案1】：

最快的选项不仅取决于 DataFrame 的长度（在这种情况下，大约 13M 行），还取决于组的数量。下面是 perfplots，它们比较了在每组中找到最大值的多种方法：

如果只有少数（大）组，using_idxmax 可能是最快的选择：

如果有很多（小）组并且 DataFrame 不太大，using_sort_drop 可能是最快的选择：

但是请记住，虽然 using_sort_drop、using_sort 和 using_rank 开始看起来非常快，但随着 N = len(df) 的增加，它们相对于其他选项的速度会很快消失。 对于足够大的N，using_idxmax 成为最快的选择，即使有很多组。

using_sort_drop、using_sort 和 using_rank 对 DataFrame（或 DataFrame 中的组）进行排序。排序平均为O(N * log(N))，而其他方法使用O(N) 操作。这就是为什么像 using_idxmax 这样的方法在非常大的 DataFrame 上胜过 using_sort_drop 的原因。

请注意，基准测试结果可能因多种原因而有所不同，包括机器规格、操作系统和软件版本。因此，在您自己的机器上运行基准测试并使用适合您情况的测试数据非常重要。

基于上面的性能图，using_sort_drop可能是对于 13M 行的 DataFrame 值得考虑的选项，尤其是在它有很多（小）组的情况下。否则，我会怀疑 using_idxmax 是最快的选择 - 但同样，检查机器上的基准非常重要。

---

这是我用来制作perfplots的设置：

import numpy as np
import pandas as pd 
import perfplot

def make_df(N):
    # lots of small groups
    df = pd.DataFrame(np.random.randint(N//10+1, size=(N, 2)), columns=['Id','delta'])
    # few large groups
    # df = pd.DataFrame(np.random.randint(10, size=(N, 2)), columns=['Id','delta'])
    return df


def using_idxmax(df):
    return df.loc[df.groupby("Id")['delta'].idxmax()]

def max_mask(s):
    i = np.asarray(s).argmax()
    result = [False]*len(s)
    result[i] = True
    return result

def using_custom_mask(df):
    mask = df.groupby("Id")['delta'].transform(max_mask)
    return df.loc[mask]

def using_isin(df):
    idx = df.groupby("Id")['delta'].idxmax()
    mask = df.index.isin(idx)
    return df.loc[mask]

def using_sort(df):
    df = df.sort_values(by=['delta'], ascending=False, kind='mergesort')
    return df.groupby('Id', as_index=False).first()

def using_rank(df):
    mask = (df.groupby('Id')['delta'].rank(method='first', ascending=False) == 1)
    return df.loc[mask]

def using_sort_drop(df):
    # Thanks to jezrael
    # https://***.com/questions/50381064/select-the-max-row-per-group-pandas-performance-issue/50389889?noredirect=1#comment87795818_50389889
    return df.sort_values(by=['delta'], ascending=False, kind='mergesort').drop_duplicates('Id')

def using_apply(df):
    selected_idx = df.groupby("Id").apply(lambda df: df.delta.argmax())
    return df.loc[selected_idx]

def check(df1, df2):
    df1 = df1.sort_values(by=['Id','delta'], kind='mergesort').reset_index(drop=True)
    df2 = df2.sort_values(by=['Id','delta'], kind='mergesort').reset_index(drop=True)
    return df1.equals(df2)

perfplot.show(
    setup=make_df,
    kernels=[using_idxmax, using_custom_mask, using_isin, using_sort, 
             using_rank, using_apply, using_sort_drop],
    n_range=[2**k for k in range(2, 20)],
    logx=True,
    logy=True,
    xlabel='len(df)',
    repeat=75,
    equality_check=check)

---

另一种基准测试方法是使用IPython %timeit：

In [55]:  df = make_df(2**20)

In [56]: %timeit using_sort_drop(df)
1 loop, best of 3: 403 ms per loop

In [57]: %timeit using_rank(df)
1 loop, best of 3: 1.04 s per loop

In [58]: %timeit using_idxmax(df)
1 loop, best of 3: 15.8 s per loop

【讨论】：

我曾经意识到最快的是df = df.sort_values(by=['delta'], ascending=False).drop_duplicates('Id')，可以添加计时吗？

当然。现在运行基准测试；需要一些时间。

@jezrael：感谢您的建议。你的方法（我在上面称为using_sort_drop）对于中等大小的DataFrames确实更快，特别是如果有很多小组。但对于非常大的 DataFrame，using_idxmax 会更快。

“中等大小”取决于组的数量。对于很少的组，using_sort_drop 在 len(df) = 100K 左右之前是有利的。但是如果有很多小团体，它可能仍然是 len(df) 大约 100M 的赢家。因此，这在很大程度上取决于 DataFrame 的性质。最好安装 perfplot 和/或 IPython 并进行一些基准测试。

唷，你真的知道如何把它拿出来！打算收藏这个，谢谢。

【参考方案2】：

使用 Numba 的 jit

from numba import njit
import numpy as np

@njit
def nidxmax(bins, k, weights):
    out = np.zeros(k, np.int64)
    trk = np.zeros(k)
    for i, w in enumerate(weights - (weights.min() - 1)):
        b = bins[i]
        if w > trk[b]:
            trk[b] = w
            out[b] = i
    return np.sort(out)

def with_numba_idxmax(df):
    f, u = pd.factorize(df.Id)
    return df.iloc[nidxmax(f, len(u), df.delta.values)]

---

借用@unutbu

def make_df(N):
    # lots of small groups
    df = pd.DataFrame(np.random.randint(N//10+1, size=(N, 2)), columns=['Id','delta'])
    # few large groups
    # df = pd.DataFrame(np.random.randint(10, size=(N, 2)), columns=['Id','delta'])
    return df

---

总理jit

with_numba_idxmax(make_df(10));

---

测试

df = make_df(2**20)


%timeit with_numba_idxmax(df)
%timeit using_sort_drop(df)

47.4 ms ± 99.8 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
194 ms ± 451 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

【讨论】：

非常感谢@piRSqured，您的代码使我的代码运行速度提高了 5 倍以上