优化多组groupby操作中公共分组列的处理方式

原始数据与场景 #

现有如下DataFrame：

In [1]: df
Out[1]:
    A   B   X1  X2  Y
0   a1  b1  p   u   1.2
1   a1  b1  p   v   0.0
2   a1  b2  q   v   1.0
3   a2  b2  p   u   2.0
4   a2  b2  q   q   5.0
5   a2  b1  q   v   7.0

需频繁执行基于[A,B,X1]和[A,B,X2]的分组聚合操作，示例代码如下：

# 按[A,B,X1]分组，计算Y的总和
df.groupby(['A','B','X1'])['Y'].sum()

# 按[A,B,X2]分组，计算Y的总和
df.groupby(['A','B','X2'])['Y'].sum()

用户猜测与问题 #

• 猜测：每次执行groupby时，都会重新计算分组键的唯一组合，比如第一次计算[A,B,X1]的唯一组合，第二次计算[A,B,X2]的唯一组合，造成重复计算的时间消耗。
• 问题1：上述猜测是否正确？
• 问题2：若猜测正确，能否利用所有分组操作中的公共键[A,B]进行优化？

上下文补充 #

• [A,B]的唯一组合数远多于X1或X2的唯一值数量；
• 当前数据各列无空值，但需保留处理分组键空值的灵活性；
• X类列数量及聚合操作类型较多，优化需求迫切。

用户曾设想过一种理想但不可行的方案：

# 不可行的理想方案
grp = df.groupby(['A','B'])
grp[['X1','Y']].groupby('X1').sum()
# 期望输出：
# A   B   X1
# a1  b1  p     1.2
#     b2  q     1.0
# a2  b1  q     7.0
#     b2  p     2.0
#         q     5.0
# Name: Y, dtype: float64

问题解答 #

问题1：猜测是否正确？ #

正确。Pandas的groupby方法在每次调用时，都会重新解析整个分组键列，计算其唯一组合并生成分组映射关系。即使分组键包含公共部分（如这里的[A,B]），也不会自动复用之前的分组结果，每次都会从头计算完整分组键的唯一组合，确实会产生重复计算的额外开销。

问题2：如何利用公共分组键[A,B]优化？ #

可以通过以下几种方式实现优化，避免重复计算[A,B]的分组：

方法1：先按[A,B]分组，子组内二次聚合

利用GroupBy.apply在每个[A,B]子组内执行针对X列的分组聚合，[A,B]的分组仅计算一次，后续X列的分组在规模更小的子组内进行，开销更低：

# 预计算[A,B]分组
ab_grp = df.groupby(['A','B'])

# 针对X1的聚合操作
result_x1 = ab_grp.apply(lambda sub_df: sub_df.groupby('X1')['Y'].sum())
# 调整索引格式，与直接groupby的输出对齐
result_x1 = result_x1.droplevel(2).rename_axis(['A','B','X1'])

# 针对X2的聚合操作同理
result_x2 = ab_grp.apply(lambda sub_df: sub_df.groupby('X2')['Y'].sum())
result_x2 = result_x2.droplevel(2).rename_axis(['A','B','X2'])

执行后result_x1的输出与直接调用df.groupby(['A','B','X1'])['Y'].sum()完全一致，且大幅减少了重复计算的开销。

方法2：预计算[A,B]分组ID，结合X列分组

先为每行分配[A,B]分组的唯一ID，后续将该ID与X列组合作为分组键，复用预计算的ID避免重复解析[A,B]组合：

# 预计算[A,B]分组的唯一ID
df['AB_group_id'] = df.groupby(['A','B']).ngroup()

# 按ID+X1分组计算Y的和
result_x1 = df.groupby(['AB_group_id','X1'])['Y'].sum()
# 将ID映射回[A,B]标签，调整索引格式
ab_labels = df[['A','B','AB_group_id']].drop_duplicates().set_index('AB_group_id')
result_x1 = result_x1.join(ab_labels).set_index(['A','B'], append=True).reorder_levels(['A','B','X1'])

# 针对X2的聚合操作同理
result_x2 = df.groupby(['AB_group_id','X2'])['Y'].sum()
result_x2 = result_x2.join(ab_labels).set_index(['A','B'], append=True).reorder_levels(['A','B','X2'])

这种方法适合需要多次针对不同X列执行聚合的场景，预计算的ID可反复使用，减少分组键解析的开销。

方法3：结合transform实现特定聚合（适合保留原行结构场景）

如果聚合操作是求和、均值这类常见类型，可使用GroupBy.transform在子组内计算，性能更优，但灵活性稍差：

# 计算每个[A,B,X1]分组下Y的和，保留原DataFrame行结构
df['Y_sum_x1'] = ab_grp['Y'].transform(lambda x: x.groupby(df['X1']).sum())

该方式会在原DataFrame中添加结果列，适合需要保留原始行数据的场景。

内容的提问来源于stack exchange，提问作者Anirban Chakraborty