如何自动将pandas dataframe分割为多个数据块？

核心实现规则 #

• 记录数阈值默认设为200000，低于阈值直接走单线程处理逻辑
• 超过阈值后按自定义线程数拆分DF，拆分逻辑完全匹配要求：优先均分，剩余的少量记录分配到最后一个/多个拆分块中
• 所有线程处理完成后自动拼接所有结果块为完整DataFrame

完整代码示例 #

import pandas as pd
from concurrent.futures import ThreadPoolExecutor

# 自定义单数据块处理逻辑，替换为你实际的计算流程
def process_chunk(chunk: pd.DataFrame) -> pd.DataFrame:
    # 示例：假设你的处理逻辑是新增一列计算结果
    # chunk['calc_result'] = chunk['col1'] + chunk['col2']
    return chunk

# 主处理流程
if __name__ == '__main__':
    # 可配置参数
    RECORD_THRESHOLD = 200000  # 启动多线程的记录数阈值
    THREAD_COUNT = 2  # 自定义线程数

    # 读取带分隔符的文件，sep替换为你的实际分隔符
    df = pd.read_csv('your_input_file.csv', sep='\t')
    total_rows = len(df)

    # 低于阈值走单线程逻辑
    if total_rows <= RECORD_THRESHOLD:
        result_df = process_chunk(df)
    else:
        # 计算拆分参数
        base_chunk_size = total_rows // THREAD_COUNT
        remainder_rows = total_rows % THREAD_COUNT

        # 拆分DataFrame为多个块
        chunks = []
        start_idx = 0
        for i in range(THREAD_COUNT):
            # 最后remainder_rows个块各多分配1条记录，匹配要求的拆分规则
            current_size = base_chunk_size + 1 if i >= (THREAD_COUNT - remainder_rows) else base_chunk_size
            end_idx = start_idx + current_size
            chunks.append(df.iloc[start_idx:end_idx].copy())
            start_idx = end_idx
        
        # 多线程并行处理
        with ThreadPoolExecutor(max_workers=THREAD_COUNT) as executor:
            processed_chunks = list(executor.map(process_chunk, chunks))
        
        # 拼接所有处理后的块
        result_df = pd.concat(processed_chunks, ignore_index=True)
    
    # 后续处理：输出result_df到文件等
    # result_df.to_csv('output_file.csv', index=False)

规则验证示例 #

• 输入200001条记录、2线程：base_chunk_size=100000，remainder_rows=1，第一个块分配100000条，第二个块分配100001条，完全匹配要求
• 输入1000000条记录、2线程：base_chunk_size=500000，remainder_rows=0，两个块各分配500000条，完全匹配要求

注意事项 #

• 如果你的计算逻辑属于CPU密集型，建议替换ThreadPoolExecutor为ProcessPoolExecutor，规避Python GIL带来的性能损耗
• 若输入文件极大，可以直接用pd.read_csv(chunksize=xxx)参数在读取阶段拆分，无需全量加载到内存后再拆分，降低内存占用

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