百万行Pandas DataFrame应用函数触发内存错误的优化求助

你当前的核心问题是用apply逐行处理百万级数据，这不仅效率极低，还会因频繁的单行数据库查询和内存碎片触发内存错误。下面是针对性的重构方案，完全避开逐行操作，用批量处理大幅提升性能并降低内存占用：

原代码问题分析 #

• apply逐行遍历本质是Python循环，对百万级数据的效率比Pandas矢量化操作低几个数量级
• 每行都发起数据库查询（Tgroup.query.filter_by().first()和Trax.query），百万次数据库IO直接拖垮性能
• 逐行修改DataFrame行对象，会产生大量内存碎片，加剧内存压力

优化后的完整代码 #

1. 主解析函数重构 #

def parse(self):
    # 可选：用chunksize分块处理，进一步降低内存占用（推荐百万级数据使用）
    chunk_size = 100000
    for chunk in pd.read_csv(file_name, chunksize=chunk_size):
        self._process_chunk(chunk)
    # 如果内存足够，也可以直接加载整个DataFrame
    # file_data_df = pd.read_csv(file_name)
    # self._process_chunk(file_data_df)

def _process_chunk(self, df):
    # 步骤1：批量处理日期，生成year、month字段
    df['period_date'] = pd.to_datetime(df['period_date'], errors='coerce')
    df['year'] = df['period_date'].dt.year
    df['month'] = df['period_date'].dt.month
    
    # 标记并记录无效日期
    invalid_rows = df['period_date'].isna()
    if invalid_rows.any():
        logging.error(f"发现{invalid_rows.sum()}条无效日期数据")
    
    # 步骤2：生成唯一分组键，批量查询已有Tgroup
    df['client_id'] = self.client_id
    unique_groups = df[['client_id', 'source', 'year', 'month']].drop_duplicates()
    group_tuples = list(unique_groups.itertuples(index=False, name=None))
    
    # 批量查询已有分组
    existing_groups = Tgroup.query.filter(
        (Tgroup.client_id, Tgroup.source, Tgroup.year, Tgroup.month).in_(group_tuples)
    ).all()
    group_id_map = {(g.client_id, g.source, g.year, g.month): g.id for g in existing_groups}
    
    # 步骤3：批量创建缺失的Tgroup
    missing_groups = [gt for gt in group_tuples if gt not in group_id_map]
    if missing_groups:
        new_tgroups = [
            Tgroup(client_id=gt[0], source=gt[1], year=gt[2], month=gt[3]) 
            for gt in missing_groups
        ]
        db.session.bulk_save_objects(new_tgroups)
        db.session.commit()
        # 更新映射字典
        for gt, tg in zip(missing_groups, new_tgroups):
            group_id_map[gt] = tg.id
    
    # 步骤4：批量生成group_id字段
    df['group_id'] = df.apply(
        lambda row: group_id_map.get((row['client_id'], row['source'], row['year'], row['month'])),
        axis=1
    )
    
    # 步骤5：批量查询最大revision，生成revision字段
    max_rev_query = Trax.query.with_entities(
        Trax.client_id, Trax.source, Trax.year, Trax.month, func.max(Trax.revision).label('max_rev')
    ).filter(
        (Trax.client_id, Trax.source, Trax.year, Trax.month).in_(group_tuples)
    ).group_by(
        Trax.client_id, Trax.source, Trax.year, Trax.month
    )
    max_rev_results = max_rev_query.all()
    revision_map = {(r.client_id, r.source, r.year, r.month): r.max_rev + 1 for r in max_rev_results}
    
    # 补充未找到的分组，默认revision=1
    for gt in group_tuples:
        if gt not in revision_map:
            revision_map[gt] = 1
    
    df['revision'] = df.apply(
        lambda row: revision_map.get((row['client_id'], row['source'], row['year'], row['month']), 1),
        axis=1
    )
    
    # 步骤6：批量入库（根据你的数据库操作逻辑调整）
    # 示例：用bulk操作批量写入Trax表
    # ... 你的批量入库代码 ...

关键优化点说明 #

1. 矢量化日期处理：用pd.to_datetime和dt.year/dt.month批量生成日期字段，比逐行判断快100+倍
2. 批量数据库查询：把百万次单行查询改成几次批量查询，彻底解决数据库IO瓶颈
3. 批量插入对象：用bulk_save_objects批量创建Tgroup，减少数据库会话的flush/commit次数
4. 分块处理：通过chunksize把大DataFrame拆成小批次，降低单批次内存占用，避免内存错误
5. 字典映射替代逐行查询：用字典存储分组ID和revision的映射，O(1)时间查询，比每行查数据库高效无数倍

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