python - Pandas dataframe groupby + apply + new column 很慢

Question

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

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

自定义函数的工作原理如下:根据 if else 条件，新列将填充 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)