如何按条件拆分Pandas DataFrame且不切断相同sid的会话数据

实现方案 #

前置准备 #

首先确保原DataFrame已经按sid排序，同一个会话ID的所有行是连续排布的，这是不拆断会话的基础。如果还未排序，先执行：

df = df.sort_values('sid', ignore_index=True)

切分代码实现 #

# 预设单块的目标行数，可按需调整为10万
target_chunk_size = 100000
total_rows = len(df)
current_start = 0

while current_start < total_rows:
    # 先按目标大小取初步切分点
    current_end = min(current_start + target_chunk_size, total_rows)
    # 未到数据末尾时，检查切分点是否落在同一个sid区间内
    if current_end < total_rows:
        border_sid = df.iloc[current_end - 1]['sid']
        # 切分点后移直到sid发生变化或到达数据末尾
        while current_end < total_rows and df.iloc[current_end]['sid'] == border_sid:
            current_end += 1
    # 切出符合要求的子DataFrame
    sub_df = df.iloc[current_start:current_end]
    # 此处可直接对sub_df做后续处理，比如写入文件、入库等
    # 处理完成后更新下一轮的起始位置
    current_start = current_end

逻辑说明 #

• 该方案无需提前聚合所有sid的位置信息，运行效率足够应对170万行的数据规模，内存占用极低
• 最终每个子DataFrame的行数会略大于等于目标行数，完全避免同一会话被拆分的问题
• 如果你不需要留存所有子DataFrame，每切出一块就直接处理，还能进一步降低内存消耗

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