Pandas 数据框 groupby + 应用 + 新列很慢

Question

我有一个 Pandas 数据框。我使用groupBy（在 1 列上）+apply 组合向数据框添加新列。 apply 调用带有参数的自定义函数。完整的调用如下所示：

df = df.groupby('id').apply(lambda x: customFunction(x,'searchString'))

自定义函数的工作方式如下：基于ifelse 条件，新列填充1 或0。然后返回该组。有点笼统，自定义函数长这样：

def customFunction(group,searchString):
    #print(group.iloc[[0]]['id'].values[0])
    if len(group[(group['name'] == searchString)) > 0:
        group['newColumn'] = 1
    else:
        group['newColumn'] = 0
    return group

我的问题是脚本运行时间相对较长，尽管我并没有处理太多数据。这些是我的数据统计： 数据框有 3130 行和 49 列。 groupBy 生成 1499 个单独的组。

如果我在customFunction 中输出一些调试文本，我观察到每个组的实际迭代相当快，但是在最后需要更多秒（比迭代本身更长）直到groupBy实际上已经完成。我认为这与重新索引或重新分配新列中的新数据有关。

我现在的问题是：

• 为什么groupBy + apply 需要这么长时间？为什么实际迭代已经完成的部分，需要这么长时间？
• 如何避免这个瓶颈？如何改进我的代码（见上文）以更快地执行？
• 更一般地说：如何最有效地实现“按特定列分组，然后根据条件添加新列”模式？也许一种方法是创建一个单独的数据结构而不需要返回组。然后，在一个单独的步骤中，新计算的数据结构可以与原始数据帧连接。但是，我不太确定这是否真的会表现得更好。

应该避免读取返回组，因为它需要很长时间，但我认为在我的情况下这是必要的，因为我在customFunction 中明确生成了新数据，这需要返回数据。

Answer 1

这是另一个更有效（针对这种特殊情况）的解决方案，没有groupby

>> searchString = 'searchString'
>> df = pd.DataFrame({'id': np.random.choice(1000, 1000000)})
>> df['name'] = random_names  # 1000000 random strings of len 10
>> df.loc[np.random.choice(1000000, 1000, replace=False), 'name'] = searchString
>>
>> def solution_0(x):
>>   x = x.groupby('id').apply(lambda g: customFunction(g, searchString))
>>
>> def solution_1(x):
>>   x['newColumn'] = x.groupby('id')['name'].transform(lambda g: g.eq(searchString).any().astype(int))
>>
>> def solution_2(x):
>>   x['newColumn'] = 0
>>   x.loc[x['id'].isin(x.loc[x['name'] == searchString, 'id']), 'newColumn'] = 1
>> 
>> %timeit solution_0(df)
3.4 s ± 125 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>> %timeit solution_1(df)
1.47 s ± 56.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>> %timeit solution_2(df)
129 ms ± 4.33 ms per loop (mean ± std. dev. of 7 runs, 10 loops each) 

Answer 2

df.groupby(...).apply(...) 没有完全矢量化，因为在底层它是一个 for .. loop，它将对每个组应用指定的函数（在你的情况下，它将被执行 1499+1 次）。

See Notes in the docs describing why Pandas apply will call func twice for the first group:

在当前的实现中，apply 在第一次调用 func 两次 组来决定它是否可以采用快速或慢速代码路径。这个可以 如果 func 有副作用，则会导致意外行为，因为它们会 第一组生效两次。

建议首先使用矢量化函数寻找解决方案，如果无法使用.apply() 作为最后的手段。

IIUC 您可以使用以下矢量化方法：

In [43]: df
Out[43]:
   id name
0   1  aaa
1   1  bbb
2   1  aaa
3   2  ccc
4   2  bbb
5   2  ccc
6   3  aaa

In [44]: searchString = 'aaa'

In [45]: df['newColumn'] = df.groupby('id')['name'] \
                             .transform(lambda x: x.eq(searchString).any().astype(int))

In [46]: df
Out[46]:
   id name  newColumn
0   1  aaa          1
1   1  bbb          1
2   1  aaa          1
3   2  ccc          0
4   2  bbb          0
5   2  ccc          0
6   3  aaa          1

时间为 70.000 行 DF：

In [56]: df = pd.concat([df] * 10**4, ignore_index=True)

In [57]: df.shape
Out[57]: (70000, 2)

In [58]: %timeit df.groupby('id').apply(lambda x: customFunction(x,searchString))
10 loops, best of 3: 92.4 ms per loop

In [59]: %timeit df.groupby('id')['name'].transform(lambda x: x.eq(searchString).any().astype(int))
10 loops, best of 3: 53.5 ms per loop