Python pandas groupby后apply函数的逐行与向量化性能差异

当你使用 pandas 的 groupby 后接 apply 函数时，是使用逐行遍历还是向量化操作，会带来巨大的性能差异。本文将直接展示这两种写法，并指导你如何选择。

理解两种方式的核心区别

逐行处理 (逐行 apply)：apply 函数接收的参数是一个分组后的子 DataFrame 或 Series，函数内部使用 iterrows、apply(axis=1) 或直接循环来处理每一行。这是一种串行过程，速度慢。

向量化处理 (分组向量化)：apply 函数接收参数后，内部直接使用 pandas 或 NumPy 的内置向量化函数（如 .sum(), .mean(), .std()）来处理整个数据块，利用底层优化的 C 代码，速度极快。

核心区别在于：是否在函数内部对“每一行”进行了显式循环。

动手实战：创建示例数据

首先，导入 pandas 库，并创建一个用于演示的 DataFrame。

import pandas as pd
import numpy as np

# 创建示例数据
data = {
    '部门': ['销售', '销售', '技术', '技术', '技术', '市场', '市场'],
    '员工': ['张三', '李四', '王五', '赵六', '钱七', '孙八', '周九'],
    '工资': [10000, 12000, 15000, 18000, 20000, 9000, 11000],
    '奖金': [1000, 1200, 2000, 2500, 3000, 800, 1000]
}
df = pd.DataFrame(data)
print(df)

输出将是一个包含 7 行数据的表格。

方法一：逐行处理（性能差）

目标：计算每个部门中，每位员工的工资加上奖金后，与该部门平均总薪资（工资+奖金）的偏差。

这个需求如果使用逐行 apply，逻辑如下：

1. 定义一个逐行处理函数。
2. 调用 groupby 后，使用 apply 并传入该函数。

def row_by_row_process(group):
    # group 是按部门分组后的一个子DataFrame
    # 1. 计算这个分组的平均总薪资（向量化操作，这步很快）
    avg_total = (group['工资'] + group['奖金']).mean()

    # 2. 对组内的每一行进行循环计算
    results = []
    for index, row in group.iterrows(): # 显式逐行循环
        total_salary = row['工资'] + row['奖金']
        deviation = total_salary - avg_total
        results.append(deviation)

    # 3. 返回一个Series，其索引需要与原分组索引对齐
    return pd.Series(results, index=group.index)

# 使用 apply 调用逐行函数
df['逐行偏差'] = df.groupby('部门').apply(row_by_row_process)
print(df[['部门', '员工', '逐行偏差']])

关键点：row_by_row_process 函数内部使用了 for index, row in group.iterrows():，这是逐行处理的典型标志。对于大数据集，iterrows 极其缓慢。

方法二：向量化处理（性能好）

目标：完成与上一节完全相同的计算。

向量化处理的思路是：在函数内部，避免循环，直接利用 pandas 对整个分组数据进行操作。

def vectorized_process(group):
    # group 是按部门分组后的一个子DataFrame
    # 所有计算都基于整个 group 进行向量化操作
    total_salary_series = group['工资'] + group['奖金']
    avg_total = total_salary_series.mean()
    # 直接对Series进行整体减法运算
    deviation_series = total_salary_series - avg_total
    return deviation_series

# 使用 apply 调用向量化函数
df['向量化偏差'] = df.groupby('部门').apply(vectorized_process)
print(df[['部门', '员工', '向量化偏差']])

关键点：vectorized_process 函数内部没有 for 循环或 iterrows。所有操作（加法、求均值、减法）都是针对整个 Series 或 DataFrame 列进行的，这就是向量化。

性能对比：为什么向量化更快？

让我们通过一个简单的计时测试来感受差异。

生成一个更大的测试数据集：

# 生成 100 个分组，每组 100 行，共 10000 行数据
np.random.seed(42)
large_df = pd.DataFrame({
    'Group': np.repeat(range(100), 100),
    'Value': np.random.randn(10000)
})

定义并测试两种函数：

# 1. 逐行函数
def slow_mean_diff(group):
    results = []
    for i, row in group.iterrows():
        diff = row['Value'] - group['Value'].mean()
        results.append(diff)
    return pd.Series(results, index=group.index)

# 2. 向量化函数
def fast_mean_diff(group):
    return group['Value'] - group['Value'].mean()

# 计时测试
import timeit

# 测试逐行 apply
time_slow = timeit.timeit(
    lambda: large_df.groupby('Group').apply(slow_mean_diff),
    number=5
)
print(f"逐行处理 (5次平均): {time_slow/5:.4f} 秒")

# 测试向量化 apply
time_fast = timeit.timeit(
    lambda: large_df.groupby('Group').apply(fast_mean_diff),
    number=5
)
print(f"向量化处理 (5次平均): {time_fast/5:.4f} 秒")

# 计算加速比
speedup = time_slow / time_fast
print(f"向量化处理比逐行处理快约 {speedup:.1f} 倍")

典型结果（实际结果因机器而异）：

• 逐行处理可能耗时数秒。
• 向量化处理可能仅需几毫秒。
• 向量化通常比逐行 apply 快 10 到 100 倍以上。

性能差异源于 iterrows 需要将每行数据转换为 Series 对象，这个过程有巨大的 Python 对象开销。而向量化操作直接在底层 C 数组上运算，没有这种开销。

结论与行动指南

何时使用 apply：

当你的分组操作逻辑极其复杂，无法用现有的 pandas/NumPy 向量化函数（如 sum, mean, transform, agg）或简单的算术运算组合完成时，才考虑使用 apply。

使用 apply 时的最佳实践：

1. 首先尝试向量化：在 apply 的函数内部，检查是否存在 for 循环、iterrows()、itertuples() 或任何显式行迭代。如果存在，尝试重写为向量化版本。
2. 如果必须逐行：对于确实需要逐行逻辑的情况（例如，调用外部 API 或执行复杂的字符串解析），请确保：
   • 明确使用 axis=1：如果是在整个 DataFrame 上调用 apply(func, axis=1)，那么函数 func 接收的是每一行。
   • 避免在分组 apply 中嵌套行循环：如前面例子所示，这是最慢的组合。
   • 考虑其他选项：评估是否可以用 transform 或 agg 替代 apply。transform 需要返回与组大小相同的结果，常用于给原 DataFrame 添加新列；agg 则更常用于聚合到组级别。

最终建议：

优先使用 groupby() + 向量化聚合函数（如 sum, mean）或 transform。仅在向量化无法表达复杂逻辑时，才使用 apply，并在其函数内部全力避免逐行循环。遵循此原则，你的 pandas 代码性能将得到数量级的提升。现在，你可以根据数据情况，选择最适合的分组计算方法了。