带条件的Groupby最佳实践：无需Merge实现及内存优化方案

需要对DataFrame执行带条件的Groupby操作，并将结果映射回原DataFrame。其中特征COL_COND的取值为1或0，需要汇总的特征为AMOUNT。

当前实现方式是通过两次Groupby操作生成Pandas Series，再通过Merge将结果合并回原DataFrame，代码如下：

import pandas as pd

df = pd.DataFrame({'ID':[1,1,2,2,3,3,3,4,5,5,6],
            'COL_COND':[1,0,1,0,1,0,1,0,1,0,0],
            'AMOUNT':[5, 80,100, 50, 100, 100, 20, 1, 51, 11, 12]})

series1 = df[df.COL_COND==1].groupby('ID')['AMOUNT'].sum().rename('sum_amount_1')
series0 = df[df.COL_COND==0].groupby('ID')['AMOUNT'].sum().rename('sum_amount_0')

df = df.merge(series1.to_frame().reset_index(), on='ID', how='left')\
.merge(series0.to_frame().reset_index(), on='ID', how='left')

print(df)

执行结果：

ID  COL_COND  AMOUNT  sum_amount_1  sum_amount_0
0    1         1       5     5.0000000            80
1    1         0      80     5.0000000            80
2    2         1     100   100.0000000            50
3    2         0      50   100.0000000            50
4    3         1     100   120.0000000           100
5    3         0     100   120.0000000           100
6    3         1      20   120.0000000           100
7    4         0       1           NaN             1
8    5         1      51    51.0000000            11
9    5         0      11    51.0000000            11
10   6         0      12           NaN            12

请问能否不使用Merge完成该操作？若存在内存占用问题，最合理的实现方法是什么？

一、不使用Merge的实现方法 #

可以通过分组变换（transform）或透视表+映射两种方式实现，无需Merge操作：

方法1：transform结合条件判断 #

直接用groupby.transform计算分组内的条件求和，一步生成目标列：

import pandas as pd

df = pd.DataFrame({'ID':[1,1,2,2,3,3,3,4,5,5,6],
            'COL_COND':[1,0,1,0,1,0,1,0,1,0,0],
            'AMOUNT':[5, 80,100, 50, 100, 100, 20, 1, 51, 11, 12]})

# 计算COL_COND=1时的分组求和
df['sum_amount_1'] = df.groupby('ID')['AMOUNT'].transform(lambda x: x[df.loc[x.index, 'COL_COND'] == 1].sum())
# 计算COL_COND=0时的分组求和
df['sum_amount_0'] = df.groupby('ID')['AMOUNT'].transform(lambda x: x[df.loc[x.index, 'COL_COND'] == 0].sum())

# 将无匹配的0值转为NaN，与原结果对齐
df['sum_amount_1'] = df['sum_amount_1'].replace(0, pd.NA)
df['sum_amount_0'] = df['sum_amount_0'].replace(0, pd.NA)

print(df)

方法2：透视表+map映射 #

先通过透视表生成ID对应的条件求和结果，再用map直接映射到原DataFrame：

import pandas as pd

df = pd.DataFrame({'ID':[1,1,2,2,3,3,3,4,5,5,6],
            'COL_COND':[1,0,1,0,1,0,1,0,1,0,0],
            'AMOUNT':[5, 80,100, 50, 100, 100, 20, 1, 51, 11, 12]})

# 生成透视表，按ID分组，COL_COND为列，聚合AMOUNT的求和值
pivot_df = df.pivot_table(
    index='ID', 
    columns='COL_COND', 
    values='AMOUNT', 
    aggfunc='sum'
).rename(columns={1:'sum_amount_1', 0:'sum_amount_0'})

# 用map将结果映射回原DataFrame
df['sum_amount_1'] = df['ID'].map(pivot_df['sum_amount_1'])
df['sum_amount_0'] = df['ID'].map(pivot_df['sum_amount_0'])

print(df)

两种方法均可得到与原代码一致的结果，且无需Merge操作。

二、内存占用优化方案 #

处理超大型DataFrame时，需减少中间对象创建，优先选择低内存开销的实现逻辑：

1. 向量化transform优化
   避免lambda中的索引查找，改用where函数实现向量化条件判断，提升效率并降低内存占用：

   df['sum_amount_1'] = df.groupby('ID')['AMOUNT'].transform(lambda x: x.where(df.loc[x.index, 'COL_COND'] == 1).sum())
   df['sum_amount_0'] = df.groupby('ID')['AMOUNT'].transform(lambda x: x.where(df.loc[x.index, 'COL_COND'] == 0).sum())
   

2. 分组键类型优化
   如果ID是高重复率的离散值，将其转为Categorical类型，可大幅降低分组操作的内存开销：

   df['ID'] = df['ID'].astype('category')
   

3. 分块处理（超大数据集场景）
   对于无法全量加载的TB级数据集，采用分块读取+增量聚合的方式：

   # 先增量计算全局ID的条件求和结果
   chunk_size = 10000
   sum_dict = {1: {}, 0: {}}
   
   for chunk in pd.read_csv('large_data.csv', chunksize=chunk_size):
       # 聚合当前分块COL_COND=1的分组和
       temp1 = chunk[chunk.COL_COND==1].groupby('ID')['AMOUNT'].sum()
       for idx, val in temp1.items():
           sum_dict[1][idx] = sum_dict[1].get(idx, 0) + val
       # 聚合当前分块COL_COND=0的分组和
       temp0 = chunk[chunk.COL_COND==0].groupby('ID')['AMOUNT'].sum()
       for idx, val in temp0.items():
           sum_dict[0][idx] = sum_dict[0].get(idx, 0) + val
   
   # 读取原数据并映射结果
   df = pd.read_csv('large_data.csv')
   df['sum_amount_1'] = df['ID'].map(sum_dict[1])
   df['sum_amount_0'] = df['ID'].map(sum_dict[0])
   

   该方法避免全量数据加载，适合处理超大规模数据集。

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