如何提高 pandas GroupBy 过滤操作的性能？

Question

这是我第一次提问。

我正在处理一个大型 CSV 数据集（它包含超过 1500 万行，大小超过 1.5 GB）。

我正在将数据提取加载到在 Jupyter Notebooks 中运行的 Pandas 数据帧中，以根据数据集推导出算法。我按 MAC 地址对数据进行分组，结果是 1+ 百万组。

我的算法开发的核心是运行这个操作：

pandas.core.groupby.DataFrameGroupBy.filter

运行此操作需要 3 到 5 分钟，具体取决于数据集。为了开发这个算法，我必须执行这个操作数百甚至数千次。

此操作似乎受 CPU 限制，并且仅使用我机器上可用的几个内核之一。我花了几个小时在线研究潜在的解决方案。我尝试使用numba 和dask 来加速此操作，但两次尝试都导致异常。

Numba 提供了一条消息，大意是“这不应该发生，感谢您帮助改进产品”。 Dask 似乎无法实现 DataFrameGroupBy.filter 操作。我无法确定如何重写我的代码以使用pool/map。

我正在寻找有关如何加速此操作的建议：

pandas.core.groupby.DataFrameGroupBy.filter

这是我的代码中此操作的示例。还有其他示例，所有这些示例的执行时间似乎都差不多。

import pandas as pd

def import_data(_file, _columns):
    df = pd.read_csv(_file, low_memory = False)
    df[_columns] = df[_columns].apply(pd.to_numeric, errors='coerce')
    df = df.sort_values(by=['mac', 'time'])
    # The line below takes ~3 to 5 minutes to run
    df = df.groupby(['mac']).filter(lambda x: x['latency'].count() > 1)
    return df

如何加快速度？

Answer 1

filter 与GroupBy 一起使用时通常会很慢。如果您尝试根据 GroupBy 中的条件过滤 DataFrame，更好的选择是使用 transform 或 map：

df[df.groupby('mac')['latency'].transform('count').gt(1)]

df[df['mac'].map(df.groupby('mac')['latency'].count()).gt(1)]