Pandas按Manager分组拆分DataFrame并导出多Sheet Excel问题

问题描述 #

现有如下DataFrame：

import numpy as np
import pandas as pd
from numpy.random import default_rng
rng = default_rng(100)
cdf = pd.DataFrame({'Id':[1,2,3,4,5],
                   'customer': rng.choice(list('ACD'),size=(5)),
                   'segment': rng.choice(list('PQRS'),size=(5)),
                   'manager': rng.choice(list('QWER'),size=(5)),
                   'dumma': rng.choice((1234),size=(5)),
                   'damma': rng.choice((1234),size=(5))
})

需要完成以下操作：

• a) 为每个manager生成一个Excel文件，文件内按segment列拆分为多个Sheet；
• b) 对segment值为Q、P、S的记录，校验dumma列值是否大于damma列值；
• c) 文件命名为{manager}.xlsx，而非out.xlsx；
• d) 若某个segment无记录，则无需创建对应的Sheet。

原尝试代码无法正常工作，仅显示R segment的内容，期望每个manager的Excel文件按需求生成对应Sheet：

DPM_col = "manager"
SEG_col = "segment"

for i,j in dict.fromkeys(zip(cdf[DPM_col], cdf[SEG_col])).keys():
    print("i is ", i)
    print("j is ", j)
    data_output = cdf.query(f"{DPM_col} == @i & {SEG_col} == @j")
    writer = pd.ExcelWriter('out.xlsx', engine='xlsxwriter')
    if len(data_output[data_output['segment'].isin(['Q','P','S'])])>0:
        if len(data_output[data_output['dumma'] >= data_output['damma']])>0:
            for seg, v in data_output.groupby(['segment']):
                v.to_excel(writer, sheet_name=f"POS_decline_{seg}",index=False)
            writer.save()
    else:
        for seg, v in data_output.groupby(['segment']):
            v.to_excel(writer, sheet_name=f"silent_inactive_{seg}",index=False)
        writer.save()

原代码问题分析 #

• 循环逻辑错误：遍历manager和segment的唯一组合，而非先按manager整体分组，导致每个manager的文件被多次覆盖，最终仅保留最后一次循环结果。
• 文件名固定：始终使用out.xlsx，未按manager动态命名，无法生成对应多个manager的文件。
• 校验逻辑错位：针对单条manager-segment组合判断，未覆盖该manager下所有符合条件的segment，且判断逻辑不准确。
• ExcelWriter创建时机错误：每次循环都创建同一个writer，导致之前写入的Sheet被覆盖。

修正后的代码 #

import numpy as np
import pandas as pd
from numpy.random import default_rng

# 生成原始数据
rng = default_rng(100)
cdf = pd.DataFrame({'Id':[1,2,3,4,5],
                   'customer': rng.choice(list('ACD'),size=(5)),
                   'segment': rng.choice(list('PQRS'),size=(5)),
                   'manager': rng.choice(list('QWER'),size=(5)),
                   'dumma': rng.choice((1234),size=(5)),
                   'damma': rng.choice((1234),size=(5))
})

DPM_col = "manager"
SEG_col = "segment"
target_segs = ['Q', 'P', 'S']

# 按manager分组处理
for manager, manager_data in cdf.groupby(DPM_col):
    # 创建对应manager的Excel写入器
    writer = pd.ExcelWriter(f"{manager}.xlsx", engine='xlsxwriter')
    
    # 按segment拆分当前manager下的记录
    for seg, seg_data in manager_data.groupby(SEG_col):
        if seg in target_segs:
            # 校验dumma >= damma，仅保留有效记录
            valid_data = seg_data[seg_data['dumma'] >= seg_data['damma']]
            if not valid_data.empty:
                valid_data.to_excel(writer, sheet_name=f"POS_decline_{seg}", index=False)
        else:
            # 非Q/P/S的segment直接写入
            seg_data.to_excel(writer, sheet_name=f"silent_inactive_{seg}", index=False)
    
    # 保存并关闭写入器
    writer.close()

代码说明 #

1. 按manager分组：先将整体DataFrame按manager字段分组，确保每个manager的记录被单独处理，避免文件覆盖。
2. 动态命名文件：使用f"{manager}.xlsx"生成对应manager的文件名，满足需求c。
3. 按需创建Sheet：在每个manager分组内按segment拆分，仅当segment有有效记录时创建对应Sheet，满足需求a和d。
4. 精准数据校验：针对Q/P/S类型的segment，筛选出dumma >= damma的有效记录后再写入，无有效记录则不创建Sheet，满足需求b。

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