按Group列分组拆分DataFrame至Excel工作表，禁止同组拆分

问题背景 #

需要将DataFrame拆分到多个Excel工作表，核心要求：

• 尽可能按近似3行一组拆分
• 同一Group列值的行必须处于同一工作表，绝对不能拆分

现有代码直接按固定行数拆分，会导致同Group的行被拆分到不同工作表，无法满足需求。

原始DataFrame #

现有问题代码 #

import pandas as pd
import numpy as np
import math

GROUP_LENGTH = 3
writer = pd.ExcelWriter('TestDataFrame.xlsx', engine='xlsxwriter')
number_of_chunks = math.ceil(len(test_data)/GROUP_LENGTH)
chunks = np.array_split(test_data, number_of_chunks)
sheet_number = 0
for chunk in chunks:
    chunk.to_excel(writer,sheet_name=f'3 Lines Sheet Number {sheet_number}')
    sheet_number+=1
writer.close()

问题：该代码按固定行数硬拆分，会将同Group的行拆分到不同工作表，比如Group 2的行可能被分到两个sheet。

解决方案代码 #

import pandas as pd

# 构建测试数据
test_data = pd.DataFrame({
    'Group': [1,1,2,2,2,3,4,4,4,5,5,5],
    'Amount': [247,194,901,111,781,610,591,220,481,421,430,449]
})

# 按Group分组，提取各组数据
grouped = test_data.groupby('Group')
groups = [group for _, group in grouped]

target_chunk_size = 3
chunks = []
current_chunk = []
current_row_count = 0

for group in groups:
    group_row_count = len(group)
    # 若当前chunk加当前组的行数超过目标值，且当前chunk已有内容，则先保存当前chunk
    if current_row_count + group_row_count > target_chunk_size and current_row_count != 0:
        chunks.append(pd.concat(current_chunk))
        current_chunk = [group]
        current_row_count = group_row_count
    else:
        current_chunk.append(group)
        current_row_count += group_row_count

# 加入最后剩余的chunk
if current_chunk:
    chunks.append(pd.concat(current_chunk))

# 写入Excel文件
with pd.ExcelWriter('TestDataGrouped.xlsx', engine='xlsxwriter') as writer:
    for idx, chunk in enumerate(chunks):
        chunk.to_excel(writer, sheet_name=f'Sheet {idx+1}', index=False)

最终拆分结果 #

• Sheet 1：包含Group 1和Group 2，共5行（接近目标行数，且同组未拆分）

• Sheet 2：包含Group 3和Group 4，共4行（符合同组不拆分要求）

• Sheet 3：包含Group 5，共3行（正好匹配目标行数）

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